#! /usr/bin/perl -w
##**************************************************************
##
## Copyright (C) 1990-2007, Condor Team, Computer Sciences Department,
## University of Wisconsin-Madison, WI.
## 
## Licensed under the Apache License, Version 2.0 (the "License"); you
## may not use this file except in compliance with the License.  You may
## obtain a copy of the License at
## 
##    http://www.apache.org/licenses/LICENSE-2.0
## 
## Unless required by applicable law or agreed to in writing, software
## distributed under the License is distributed on an "AS IS" BASIS,
## WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
## See the License for the specific language governing permissions and
## limitations under the License.
##
##**************************************************************

use strict;

# Update the module include path
BEGIN
{
    my $Dir = $0;
    if ( $Dir =~ /(.*)\/.*/ )
    {
	push @INC, "$1";
    }
}
use HawkeyePublish;
use HawkeyeLib;

# Setup the hawkeye stuff
my $Hawkeye;

# Setup
my %Config = (
	      MasterOwnerList => "",
	      DaemonOwnerList => "",
	      UserOwnerList => "",
	      IgnoreCmdRe => "",
	      ExtraCmdList => "",
	     );

# Do it
Init();
RunIt();

# Init logic
sub Init {
    HawkeyeLib::DoConfig( );

    $Hawkeye = HawkeyePublish->new;
    $Hawkeye->Quiet( 1 );
    $Hawkeye->AutoIndexSet( 1 );

    my $Tmp;
    $Tmp = HawkeyeLib::ReadConfig( "_master_owners", "" );
    $Config{MasterOwnerList} = $Tmp if ( $Tmp ne "" );
    $Tmp = HawkeyeLib::ReadConfig( "_daemon_owners", "" );
    $Config{DaemonOwnerList} = $Tmp if ( $Tmp ne "" );
    $Tmp = HawkeyeLib::ReadConfig( "_user_owners", "" );
    $Config{UserOwnerList} = $Tmp if ( $Tmp ne "" );
    $Tmp = HawkeyeLib::ReadConfig( "_ignore_cmd_re", "" );
    $Config{IgnoreCmdRe} = $Tmp if ( $Tmp ne "" );
    $Tmp = HawkeyeLib::ReadConfig( "_extra_cmd_list", "" );
    $Config{ExtraCmdList} = $Tmp if ( $Tmp ne "" );
}


# Do the real work here...
sub RunIt {

    # Start things off
    my $Hash = HawkeyeHash->new( \$Hawkeye, "" );

    # Data we'll gather...
    my %CondorInfo = (
		      "NumStarters"	=> 0,
		      "NumStartds"	=> 0,
		      "NumSchedds"	=> 0,
		      "NumMasters"	=> 0,
		      "NumExecs"	=> 0,
		      "NumShadows"	=> 0,
		      "NumRunaway"	=> 0,
		      "Slots"		=> "",
		      "RunawayPids"	=> { },
		     );

    # List of owners, keyed by type
    my $DefaultDaemonOwnerList = "condor,root";
    my %OwnerStrings = (
		      User	=> "*",
		      Daemon	=> $DefaultDaemonOwnerList,
		      Master	=> $DefaultDaemonOwnerList,
		     );

    # Override the defaults
    if ( $Config{UserOwnerList} ne "" )
    {
	$OwnerStrings{User} = $Config{UserOwnerList};
    }
    if ( $Config{DaemonOwnerList} ne "" )
    {
	$OwnerStrings{Daemon} = $Config{DaemonOwnerList};
	$OwnerStrings{Master} = $Config{DaemonOwnerList};
    }
    if ( $Config{MasterOwnerList} ne "" )
    {
	$OwnerStrings{Master} = $Config{MasterOwnerList};
    }

    # Build the actual lists list
    my %OwnerLists;
    foreach my $Key ( keys %OwnerStrings )
    {
	@{$OwnerLists{$Key}} = split( /,/, $OwnerStrings{$Key} );
    }

    # Table to translate "ps-command" regex to $CondorInfo key..
    my %CondorProgTable =
	(
	 condor_starter	=> {
			    Attr => "NumStarters",
			    Owner => "Daemon",
			   },
	 condor_startd	=> {
			    Attr => "NumStartds",
			    Owner => "Daemon",
			   },
	 condor_master	=> {
			    Attr => "NumMasters",
			    Owner => "Master",
			    NeverRunAway => 1,
			   },
	 condor_schedd	=> {
			    Attr => "NumSchedds",
			    Owner => "Daemon",
			   },
	 condor_exec	=> {
			    Attr => "NumExecs",
			    Owner => "User",
			   },
	 condor_shadow	=> {
			    Attr => "NumShadows",
			    Owner => "User",
			   },
	);

    my @RunAwayReList;
    push( @RunAwayReList, "condor_execute" );
    if ( $Config{ExtraCmdList} ne "" )
    {
	my $Tmp = $Config{ExtraCmdList};
	$Tmp =~ s/\//\\\//g;
	push( @RunAwayReList, split( /\s*,\s*/, $Tmp ) );
    }
    foreach my $Name ( keys %CondorProgTable )
    {
	if ( not exists $CondorProgTable{$Name}{NeverRunAway} )
	{
	    push( @RunAwayReList, $Name );
	}
    }
    my $RunAwayRe = "(" . join( "|", @RunAwayReList ) . ")";
    #print STDERR "RunAwayRe = /$RunAwayRe/\n";

    # Run a 'ps' and gather some info...
    my $Cmd = "/bin/ps -eo user,pid,ppid,comm,args";
    if ( ! open( PS, "$Cmd|" ) )
    {
	print STDERR "Warning; failed to run '$Cmd'\n";
    }
    else
    {
	my $IgnoreRe = $Config{IgnoreCmdRe};
	while ( <PS> )
	{
	    chomp;
	    my ( $User, $Pid, $Ppid, @ProgString ) = split;
	    my $ProgString = join(" ", @ProgString);

	    next if (/^USER\s+PID\s+PPID/ );
	    next if ( ( $IgnoreRe ne "" ) and
		      ( $ProgString =~ /$IgnoreRe/ ) );

	    # Starter with PPID == 1
	    if ( 1 == $Ppid )
	    {
		if ( $ProgString =~ /$RunAwayRe/ )
		{
		    $CondorInfo{NumRunaway}++;
		    $CondorInfo{RunawayPids}{$Pid} = $ProgString;
		    #print STDERR "'$ProgString' is a runaway\n";
		}
	    }

	    # Count the processes...
	    foreach my $ProgPat ( keys %CondorProgTable )
	    {
		if ( $ProgString =~ /$ProgPat/ )
		{
		    my $Attr = $CondorProgTable{$ProgPat}->{Attr};
		    my $OwnerType = $CondorProgTable{$ProgPat}->{Owner};

		    # Is it on the user match list?
		    my $UserMatch = 0;
		    foreach my $TmpUser ( @{$OwnerLists{$OwnerType}} )
		    {
			$UserMatch++ if (  ( $TmpUser eq "*" ) ||
					   ( $TmpUser eq $User )  );
		    }
		    last if ( ! $UserMatch );

		    # Update the count...
		    if ( ! exists $CondorInfo{$Attr} )
		    {
			$CondorInfo{$Attr} = 1;
		    }
		    else
		    {
			$CondorInfo{$Attr}++;
		    }
		}
	    }
	}
	close( PS );
    }

    # Now, run condor_status and gather some more info...
    if ( $CondorInfo{NumStartds} )
    {
	my %Interesting =
	    ( Name =>			{ Required => 1, },
	      Arch =>			{ Required => 1, },
	      OpSys =>			{ Required => 1, },
	      State =>			{ Required => 1, },
	      EnteredCurrentState =>	{ Required => 1, },
	      Activity =>		{ Required => 1, },
	      EnteredCurrentActivity =>	{ Required => 1, },
	      LoadAvg =>		{ Required => 1, },
	      Memory =>			{ Required => 1, },
	      LastHeardFrom =>		{ Required => 1, },
	      StarterAbilityList =>	{ Required => 0, },
	    );
	if ( open( CSTAT, "condor_status -l -direct $ENV{HOST} |" ) )
	{
	    my @Slots;
	    my %AdInfo;
	  READLINE:
	    while ( <CSTAT> )
	    {
		chomp;

		# Empty line between ads...
		if ( ! /\S+/ )
		{
		    # Verify that we have all the interesting stuff
		    next if ( scalar( keys( %AdInfo ) ) == 0 );
		    foreach my $Key ( keys %Interesting )
		    {
			if (  ( $Interesting{$Key}{Required} ) &&
					    ( ! exists $AdInfo{$Key} )  )
			{
			    print STDERR
				"Throwing out ad because '$Key' is missing\n";
			    %AdInfo = ( );
			    next READLINE;
			}
		    }

		    # Ok, now pull  it apart..
		    my $Slot = "";
		    if ( $AdInfo{Name} =~ /(slot\d+)@(.*)/ )
		    {
			$Slot = $1;
			push( @Slots, $Slot );
		    }
		    $CondorInfo{OpSys} = $AdInfo{OpSys};
		    $CondorInfo{Arch} = $AdInfo{Arch};
		    $CondorInfo{StarterAbilityList} = $AdInfo{StarterAbilityList}
			if ( exists $AdInfo{StarterAbilityList} );
		    $CondorInfo{"State$Slot"} = $AdInfo{State};
		    $CondorInfo{"Activity$Slot"} = $AdInfo{Activity};
		    $CondorInfo{"LoadAv$Slot"} = $AdInfo{LoadAvg};
		    $CondorInfo{"Mem$Slot"} = $AdInfo{Memory};
		    $CondorInfo{"ActivityTime$Slot"} =
			( $AdInfo{LastHeardFrom} - $AdInfo{EnteredCurrentActivity} );
		    %AdInfo = ( );
		    next;
		}

		# Skip the 'uniteresting' lines
		my ( $Attr, $Value ) = split( / = /, $_, 2 );
		if ( $Value =~ /^\"(.*)\"$/ )
		{
		    $Value = $1;
		}

		# De we care about it??
		next if ( ! exists $Interesting{$Attr} );
		$AdInfo{$Attr} = $Value;
	    }
	    close( CSTAT );
	    $CondorInfo{Slots} = join( " ", @Slots );
	}
	else
	{
	    print STDERR "Can't query startd\n";
	}
    }

    # Search for core files laying around
    {
	my $LogDir = `condor_config_val LOG`;
	chomp $LogDir;

	if ( ( $LogDir ne "" ) && ( -d $LogDir ) )
	{
	    if ( ! opendir( LOG, $LogDir ) )
	    {
		warn "Can't read log directory '$LogDir'";
	    }
	    else
	    {
		my @CoreFiles = grep( /^core/, readdir( LOG ) );
		$CondorInfo{CoreFilesNum} = $#CoreFiles + 1;
		my $TotalSize = 0;
		my $CoreNum = 0;
		my $Now = time( );
		foreach my $Core ( @CoreFiles )
		{
		    my $File = "$LogDir/$Core";
		    if ( my @Stat = stat( $File ) )
		    {
			my $Name = "CoreFiles_$CoreNum";
			$CondorInfo{$Name."_Name"} = $Core;
			$CondorInfo{$Name."_Size"} = $Stat[7];
			$TotalSize += ( $Stat[11] * $Stat[12] );
			$CondorInfo{$Name."_Age"} = $Now - $Stat[9];
			$CondorInfo{$Name."_TimeStamp"} = $Stat[9];
		    }
		    $CoreNum++;
		}
		$CondorInfo{CoreFilesDisk} = $TotalSize;
	    }
	    closedir( LOG );
	}
    }

    # Publish 'em all
    foreach my $Key ( keys %CondorInfo )
    {
	if ( ref( $CondorInfo{$Key} ) eq "ARRAY" )
	{
	    $Hash->Add( $Key,
			HawkeyePublish::TypeString,
			join( ",", @{$CondorInfo{$Key}} ) );
	}
	elsif ( ref( $CondorInfo{$Key} ) eq "HASH" )
	{
	    my @List = keys %{$CondorInfo{$Key}};
	    $Hash->Add( $Key,
			HawkeyePublish::TypeString,
			join( ",", @List ) );
	    foreach my $ListKey ( @List )
	    {
		$Hash->Add( $Key . "_" . $ListKey,
			    HawkeyePublish::TypeAuto,
			    $CondorInfo{$Key}{$ListKey} );
	    }
	}
	else
	{
	    $Hash->Add( $Key,
			HawkeyePublish::TypeAuto,
			$CondorInfo{$Key} );
	}
    }


    # Ok, summary is done...
    $Hash->Store( );
    $Hawkeye->Publish( );
}
