root/branch/ggpolo/umitCore/Scheduler.py @ 770

# Copyright (C) 2007 Insecure.Com LLC.
#
# Original author: Adriano Monteiro Marques
# Now maintained and updated by: Guilherme Polo <ggpolo@gmail.com>
#
# 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., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA

"""
Job scheduler
"""

import os
import sys
import time
import warnings
from tempfile import mktemp
from ConfigParser import ConfigParser

from umitCore.CronParser import CronParser
from umitCore.NmapCommand import NmapCommand
from umitCore.SendMail import send_mail
from umitCore.Paths import Path
from umitDB.xmlstore import XMLStore


try:
    umitdb_ng = Path.umitdb_ng
except:
    from umitCore.Paths import Path
    Path.set_umit_conf(join(split(__file__)[0], 'config', 'umit.conf'))

class Scheduler(object):
    """
    Schedules schemas to run.
    """

    def __init__(self, schemas_file):
        self.schemas_file = schemas_file
        self.schemas_stat = None

    def parse_schemas(self):
        """
        Parse schemas set in file.
        """
        self.schemas = [ ]

        self.schema_parser = ConfigParser()
        self.schema_parser.read(self.schemas_file)

        for sec in self.schema_parser.sections():
            self.schemas.append(SchedSchema(self.schema_parser, sec))

    
    def check_for_changes(self):
        """
        If schemas file changed since last check, reparse schemas.
        """
        try:
            new_stat = os.stat(self.schemas_file)
        except OSError: # file was being saved (probably) when this was called
            time.sleep(.1)
            new_stat = os.stat(self.schemas_file)

        new_stat = new_stat.st_mtime

        if new_stat != self.schemas_stat:
            print "Schemas file changed since last check"
            self.schemas_stat = new_stat
            self.parse_schemas()


    def _get_schemas_stat(self):
        """
        Get latest os.stat for schemas file
        """
        return self.__schemas_stat


    def _set_schemas_stat(self, stat):
        """
        Set current os.stat for schemas file
        """
        self.__schemas_stat = stat


    def _get_schemas_file(self):
        """
        Get schemas file.
        """
        return self.__schemasf


    def _set_schemas_file(self, file):
        """
        Set schemas file.
        """
        try:
            os.stat(file)
        except OSError, e:
            print e
            sys.exit(0)

        self.__schemasf = file


    # Properties
    schemas_file = property(_get_schemas_file, _set_schemas_file)
    schemas_last_stat = property(_get_schemas_stat, _set_schemas_stat)


class SchedSchema(object):
    """
    Parse scheduled schemas.
    """

    def __init__(self, schema_parser, schema_name):
        self.options = { 
            # Task execution time
            'hour':self._set_hour,
            'minute':self._set_minute,
            'month':self._set_month,
            'day':self._set_day,
            'weekday':self._set_weekday,
            # Task command
            'command':self._set_command,
            # Task options
            'enabled':self._set_enabled,
            'addtoinv':self._set_addtoinv,
            'saveto':self._set_saveto,
            'mailto':self._set_mailto
        }

        self.__schema_name = schema_name
        self.schema_parser = schema_parser
        self.last_check = None
        self.cron_parser = CronParser()
        self.setup_schema(self.schema_name)


    def job_to_run(self):
        """
        Check if there is a scheduled job for now and if there is,
        return True, otherwise False.
        """
        cur_time = time.localtime()

        # check if we should check for job or not
        if self.last_check and cur_time[1] == self.last_check[1] and \
            cur_time[2] == self.last_check[2] and \
            cur_time[3] == self.last_check[3] and \
            cur_time[4] == self.last_check[4] and \
            cur_time[6] == self.last_check[6]:

            return False # too early to run a job again!

        if cur_time[1] in self.month and cur_time[2] in self.day \
            and cur_time[3] in self.hour and cur_time[4] in self.minute \
            and cur_time[6] in self.weekday:
            
            self.last_check = cur_time
            return True # there is a job to run!


    def setup_schema(self, name):
        """
        Setup schema.
        """
        if self.schema_parser.has_section(name):
            for opt in self.schema_parser.options(name):
                if opt in self.options.keys():
                    self.options[opt](self.schema_parser.get(name, opt))


    def _get_schema_name(self):
        """
        Return schema name set.
        """
        return self.__schema_name


    def _get_last_check(self):
        """
        Return when last check happened.
        """
        return self.__lcheck


    def _set_last_check(self, when):
        """
        Set last check time.
        """
        self.__lcheck = when


    # Job time
    def _get_month(self):
        return self.__month


    def _set_month(self, month):
        """
        Parse and set month.
        """
        self.__month = self.cron_parser.parse_month(month)


    def _get_day(self): 
        return self.__day


    def _set_day(self, day):
        """
        Parse and set day.
        """
        self.__day = self.cron_parser.parse_day(day)


    def _get_weekday(self): 
        return self.__weekday


    def _set_weekday(self, weekday):
        """
        Parse and set weekday.
        """
        self.__weekday = self.cron_parser.parse_weekday(weekday)


    def _get_hour(self): 
        return self.__hour


    def _set_hour(self, hour):
        """
        Parse and set hour.
        """
        self.__hour = self.cron_parser.parse_hour(hour)


    def _get_minute(self): 
        return self.__minute


    def _set_minute(self, minute):
        """
        Parse and set minute.
        """
        self.__minute = self.cron_parser.parse_minute(minute)


    # Job Command
    def _get_command(self):
        """
        Get job command.
        """
        return self.__command


    def _set_command(self, command):
        """
        Set command for job.
        """
        self.__command = command


    # Job Options
    def _get_enabled(self):
        """
        Returns True if job should run, otherwise, False.
        """
        return self.__enabled


    def _set_enabled(self, enable):
        """
        Set if job should run or not.
        """
        self.__enabled = enable


    def _get_addtoinv(self):
        """
        Returns True if job result should be stored in Inventory, otherwise,
        False.
        """
        return self.__addtoinv


    def _set_addtoinv(self, add):
        """
        Sets job result to be added to Inventory, or not.
        """
        self.__addtoinv = add


    def _get_saveto(self):
        """
        Get file that stores job results.
        """
        return self.__savefile


    def _set_saveto(self, file):
        """
        Set a file to store job results.
        """
        self.__savefile = file
        

    def _get_mailto(self):
        """
        Get email set to receive job results.
        """
        return self.__mail


    def _set_mailto(self, mail): 
        """
        Set an email to receive job results.
        """
        self.__mail = mail


    # Properties
    month = property(_get_month, _set_month)
    day = property(_get_day, _set_day)
    weekday = property(_get_weekday, _set_weekday)
    hour = property(_get_hour, _set_hour)
    minute = property(_get_minute, _set_minute)
    command = property(_get_command, _set_command)
    enabled = property(_get_enabled, _set_enabled)
    addtoinv = property(_get_addtoinv, _set_addtoinv)
    saveto = property(_get_saveto, _set_saveto)
    mailto = property(_get_mailto, _set_mailto)

    last_check = property(_get_last_check, _set_last_check)
    schema_name = property(_get_schema_name)


def decide_ouput(file):
    """
    Choose a better or the same output path.
    """
    output = None
    err = True
    file_path = os.path.abspath(file)

    if os.path.exists(os.path.split(file_path)[0]): # path exists at least
        try: # try to open file in read mode
            f = open(file, 'r')
            f.close()
            # if we are still here, file exists, this is bad ;/
            count = 1
            orig = file
            while os.path.isfile(file):
                file = ''
                file = orig + "_%d" % count # try appending extra extension
                count += 1
            err = False
            warnings.warn("File will be saved as %s" % file, Warning)
            
        except IOError: # file didnt exist and path exists, this is good!
            err = False
            
        output = file
        
    return output


def calc_next_time():
    """
    Calculate next time to check for changes and jobs.
    """
    tt = time.localtime(time.time() + 60)
    tt = tt[:5] + (0,) + tt[6:]
    return time.mktime(tt)


running_scans = { }

def stop():
    """
    Clean temp files used by running scans.
    """    
    for running, opts in running_scans.items():
        print "Cleaning up scan \"%s\"" % opts[3]
        scan = opts[0]
        scan.close() # delete temporary files
        

def run(schema_file=None):
    """
    Run scheduler forever.
    """
    if not schema_file:
        print "Get file from umitCore.Path"
        print "Not working yet"
        sys.exit(0)

    next_time = calc_next_time()

    s = Scheduler(schema_file)

    scount = 0
    while 1: # run forever and ever ;)
        current_time = time.time()

        if current_time < next_time:
           time.sleep(next_time - current_time + .1)

        # check if time has changed by more than two minutes (clock changes)
        if abs(time.time() - next_time) > 120:
            # reset timer
            next_time = calc_next_time()

        s.check_for_changes()

        for schema in s.schemas:
            if schema.job_to_run():
                if not int(schema.enabled): # schema disabled, neeexxt!
                    continue

                print "To run:", schema.schema_name

                scan = NmapCommand(schema.command)
                scan.run_scan()

                running_scans[scount] = (scan, schema.saveto, schema.mailto,
                                         schema.schema_name, schema.addtoinv)
                scount += 1
        
        for running, opts in running_scans.items():
            scan = opts[0]
            if not scan.scan_state(): # scan finished
                print "Scan finished:", opts[3]
            
                if opts[1]: # save xml output
                    saveto = decide_ouput(opts[1])
                    f = open(saveto, 'w')
                    f.write(open(scan.get_xml_output_file(), 'r').read())
                    f.close()
                    print "Saved as:", saveto
                
                if opts[2]: # mail output to
                    recipients = opts[2].split(',')
                    print "Mail to:", recipients
                    send_mail("sendmail", recipients,
                              _("UMIT: Status Report for scheduled schema \
\"%s\"" % opts[3]),
                              _("There was a scheduled job that finished \
now: %s\nFollows an attachment of job output." % time.ctime()),
                              scan.get_xml_output_file(), 
                              
                          )
                
                if int(opts[4]): # add to inventory
                    print "Results being added to Inventory %s" % opts[3]
                    XMLStore(Path.umitdb_ng, scan.get_xml_output_file(), 
                             inventory=opts[3])
                    print "Finished insertion on Inventory %s" % opts[3]
                    
                scan.close() # delete temporary files
                scount -= 1
                del running_scans[running]

        next_time += 60


if __name__ == "__main__":
    try:
        run('schema-sample.conf')
    except KeyboardInterrupt:
        stop()