
'''
File                : proc.py 

Start Date        : 20050708 

Description        :
  This is an Eddie data collector.  It collects process information on
  a Win32 system.

  Requires Mark Hammond's win32all package.

  The following statistics are currently collected and made available to
  directives that request it (e.g., PROC):

  Stats available for processes are:
      pid                - Process ID [int]
      exe                - Full executable name [string]
      procname                - Base executable name (without path) [string]

$Id: proc.py 895 2007-12-09 07:18:02Z chris $
'''

__version__ = '$Revision: 895 $'

__copyright__ = 'Copyright (c) Chris Miles 2005'

__author__ = 'Chris Miles'

__license__ = '''
    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the Free Software
    Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
'''



# Python modules
import os

# Win32 modules
import win32process
import win32api
import win32con
import pywintypes

# Eddie modules
from eddietool.common import datacollect, log, utils


# List of interpreters - default empty - not used for Windows
interpreters = []

class procList(datacollect.DataCollect):
    """Class procList, holds a list of running processes and related information.
    """

    def __init__(self):
        apply( datacollect.DataCollect.__init__, (self,) )


    ##################################################################
    # Public, thread-safe, methods

    def procExists(self, procname):
        """Search the process list and return the number of occurrences
        of procname.
        """

        proclist = self.getList( 'proclist' )   # safely get copy of process list

        count = 0               # count number of occurrences of 'procname'
        for i in proclist:
            if i.procname == procname:
                count = count + 1

        return count


    def pidExists(self, pid):
        """Return true (1) if a process with 'pid' exists,
        or false (0) otherwise.
        """

        prochash = self.getHash( 'datahash' )   # safely get copy of process dict

        try:
            prochash[pid]
            return 1
        except KeyError:
            return 0


    def __getitem__(self, name):
        """Overload '[]' to return corresponding process object for given
        process name.

        Note that when processes are collected there may be multiple
        processes with the same name.  This dictionary only keeps a
        reference to one of those (no guarantee which one).
        """

        processes = self.getHash( 'nameHash' )  # safely get copy of process name dictionary

        try:
            r = processes[name]
        except KeyError:
            r = None            # no such process called 'name'

        return r


    def allprocs(self):
        """Return dictionary of all processes (which are dictionaries of each
        process' details).
        """

        processes = self.getHash( 'nameHash' )  # safely get copy of process name dictionary

        allprocs = {}

        for p in processes.keys():
            allprocs[p] = processes[p].procinfo()

        return allprocs


    ##################################################################
    # Private methods.  No thread safety if not using public methods.

    def collectData(self):
        """Collect process list.
        """

        self.data.datahash = {}                # dict of processes keyed by pid
        self.data.proclist = []                # list of processes
        self.data.nameHash = {}                # dict of processes keyed by process name

        procs = win32process.EnumProcesses()

        for pid in procs:
            try:
                han = win32api.OpenProcess(win32con.PROCESS_QUERY_INFORMATION|win32con.PROCESS_VM_READ, 0, pid)
            except:
                # privileges may prevent querying the process details
                han = None
            p = proc(pid, han)

            if han:
                han.Close()

            self.data.proclist.append(p)
            self.data.datahash[p.pid] = p
            self.data.nameHash[p.procname] = p

        log.log( "<proc>procList.collectData(): new proc list created", 7 )



class proc:
    """Class proc : holds information about a single process.
    """

    def __init__(self, pid, han):

        if han:
            try:
                exe = win32process.GetModuleFileNameEx(han, 0)
            except pywintypes.error, msg:
                if msg[0] != 299 : # we keep getting this error for some reason..?
                    # Error is "Only part of a ReadProcessMemory or WriteProcessMemory request was completed."
                    log.log( "<proc>proc.__init__(): failed win32process.GetModuleFileNameEx(), %s" %(msg), 5 )
                exe = "<unknown>"
        else:
            exe = "<unknown>"

        self.pid = pid                        # process ID
        self.exe = exe                        # full executable name (incl path)
        self.procname = os.path.basename(exe)        # executable name (no path)


    def procinfo(self):
        """Return process details as a dictionary.
        """

        info = {}
        info["pid"] = self.pid
        info["exe"] = self.exe
        info["procname"] = self.procname

        return info


##
## END - proc.py
##
