#!/usr/bin/python

#* Skynet - Automated "Cloud" Security Scanner                                *#
#* Copyright (C) 2014-present  Jason Frisvold <friz@godshell.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., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA *#

# Import libraries
import sys
import getopt
import ConfigParser
import time
import sqlite3
import logging
from datetime import datetime
import os, os.path
from subprocess import Popen
import re
import shlex
import tempfile

sys.path.append('../libs')
from skynettimer import SkynetTimer
from funcs import pid_exists

# Global Variables
verbose = False
programname = 'SkyNet Spawner'
version = '1.0'
configfile = 'spawner.conf'
defaultconfigfile = 'spawner.default.conf'
cfg = {}

# Set up the logger
logging.basicConfig()
logger = logging.getLogger('spawner')

###
# Main
###
def main(argv):
    try:
        opts, args = getopt.gnu_getopt(argv,
                                       "dhvVs:c:",
                                       ["debug", "help", "verbose", "config",
                                        "version", "license"])
    except getopt.GetoptError, err:
        print str(err)
        usage()
        sys.exit(2)

    for o, a in opts:
        if o in ("-v", "--verbose"):
            logger.setLevel(logging.INFO)
        elif o in ("-d", "--debug"):
            logger.setLevel(logging.DEBUG)
        elif o in ("-c", "--config"):
            global configfile
            # TODO : This should be validated
            configfile = o
        elif o in ("-V", "--version"):
            print '{0} {1}'.format(programname, version)
            sys.exit()
        elif o in ("-h", "--help"):
            usage()
            sys.exit()
        elif o in ("--license"):
            license()
            sys.exit()
        else:
            assert False, "unhandled option"

    loadconfig()

    # Main loop
    while True:
        loopstarttime = datetime.now()

        logger.info('Loop starts at {0}'.format(str(loopstarttime)))

        timingdb = open_timing_database()

        filelist = check_configdir()
        if (len(filelist) > 0):
            process_config(timingdb, filelist)

        spawnlist = check_timing(loopstarttime, timingdb)

        spawnlist = check_complete(spawnlist, timingdb)

        if (len(spawnlist) > 0):
            spawn_process(spawnlist)

        # Check for child exit status so we don't have zombies
        try:
            os.waitpid(-1, os.WNOHANG)
        except:
            pass

        logger.info('Sleeping until next loop iteration')

        # Sleep until the beginning of the next minute
        sleeptime = 60 - datetime.utcnow().second
        time.sleep(sleeptime)

def loadconfig():
    logger.info('Loading spawner configuration')

    config = ConfigParser.SafeConfigParser()

    config.read(defaultconfigfile)

    try:
        config.read(configfile)
    except:
        e = sys.exc_info()[0]
        logger.exception('Unable to load config file : {0}'.format(e))

    global cfg

    cfg['configdir']    = config.get('spawner', 'configdir')
    cfg['datadir']      = config.get('spawner', 'datadir')
    cfg['timingdb']     = config.get('spawner', 'timingdb')
    cfg['nmap_binary']  = config.get('spawner', 'nmap_binary')
    cfg['output_dir']   = config.get('spawner', 'output_dir')
    cfg['gpg_binary']   = config.get('spawner', 'gpg_binary')
    cfg['gpg_key']      = config.get('spawner', 'gpg_key')

    logger.debug('configdir = {0}'.format(cfg['configdir']))
    logger.debug('datadir = {0}'.format(cfg['datadir']))
    logger.debug('timingdb = {0}'.format(cfg['timingdb']))
    logger.debug('nmap_binary = {0}'.format(cfg['nmap_binary']))
    logger.debug('output_dir = {0}'.format(cfg['output_dir']))

def check_configdir():
    logger.info('Check Config')

    if (os.path.exists(cfg['configdir'])):
        # Look for new files

        filelist = []

        for name in os.listdir(cfg['configdir']):
            if re.match('.*\.skynet$', name):
                filelist.append(name)

        return filelist
    else:
        logger.critical('Configuration directory does not exist')
        sys.exit()

def process_config(timingdb, filelist):
    logger.info('Process Config')

    for file in filelist:
        logger.debug('Processing file - {0}'.format(file))
        incoming = open(os.path.join(cfg['configdir'], file), 'r')
        lines = incoming.readlines()
        incoming.close()

        try:
            timer = SkynetTimer()
            timer.loadfromfile(lines)
            timer.sendtodb(timingdb)
            os.unlink(os.path.join(cfg['configdir'], file))

        except:
            e = sys.exc_info()[0]
            logger.exception('Error creating timing object : {0}'.format(e))

def check_timing(loopstarttime, timingdb):
    logger.info('Check Timing')

    spawnlist = []

    try:
        timingcursor = timingdb.cursor()

        sql = '''SELECT t.timer_id, td.override_flag FROM timers AS t,
                 timer_details AS td WHERE (t.minute = ? OR t.minute = "*") AND
                 (t.hour = ? OR t.hour = "*") AND (t.day = ? OR t.day = "*")
                 AND (t.month = ? OR t.month = "*") AND t.timer_id =
                 td.timer_id'''

        timingcursor.execute(sql, (loopstarttime.minute, loopstarttime.hour,
                             loopstarttime.day, loopstarttime.month))

        rows = timingcursor.fetchall()

        if (len(rows) > 0):
            for row in rows:
                spawnlist.append((row[0], row[1]))

    except:
        e = sys.exc_info()[0]
        logger.exception('Error checking timing : {0}'.format(e))

    return spawnlist

def check_complete(spawnlist, timingdb):
    logger.info('Check for completed processes')

    try:
        timingcursor = timingdb.cursor()

        sql = 'SELECT timer_id, start_time, pid FROM spawned'

        timingcursor.execute(sql)

        rows = timingcursor.fetchall()

    except:
        e = sys.exc_info()[0]
        logger.exception('Error getting spawned process data : {0}'.format(e))

    if (len(rows) > 0):
        for row in rows:
            logger.debug('Checking id {0} started at {1}'.format(str(row[0]),
                                                                 str(row[1])))

            if pid_exists(row[2]):
                logger.debug('PID {0} Still running'.format(str(row[2])))

                try:
                    # If timer_id should be spawned, isn't allowed to override,
                    # and is still running, then remove it from the spawnlist
                    # and set overrun = 1
                    spawnlist.remove((row[0],0))
                    logger.debug('''ID {0} not allowed to override, removed from
                                    spawn list'''.format(str(row[0])))

                    sql = '''UPDATE spawned SET overtime = 1 WHERE timer_id =
                             ? AND pid = ?'''

                    timingcursor.execute(sql, (row[0], row[2],))
                except:
                    pass
            else:
                logger.debug('PID {0} Completed'.format(str(row[2])))
                end_time = int(time.mktime(time.gmtime()))

                sql = '''DELETE FROM spawned WHERE timer_id = ? AND pid = ?'''
                timingcursor.execute(sql, (row[0], row[2]))

                sql = '''INSERT INTO spawn_log (timer_id, start_time, end_time)
                         VALUES (?, ?, ?)'''
                timingcursor.execute(sql, (row[0], row[1], end_time))

    timingdb.commit()

    return spawnlist

def spawn_process(spawnlist):
    logger.info('Spawn Process')

    for (timer, override) in spawnlist:
        start_time = int(time.mktime(time.gmtime()))
        logger.debug('Spawning process for timer : {0}'.format(str(timer)))

        # First fork, if pid is zero, then it's the child
        try:
            pid = os.fork()
            if pid > 0:
                # The parent process should jump to the next loop iteration
                continue
        except OSError, e:
            sys.stderr.write('fork #1 failed: ({0}) {1}'.format(e.errno,
                                                                e.strerror))
            sys.exit(1)

        # Set the parameters for the child
        os.chdir('.')
        os.umask(0)
        os.setsid()

        # Second fork, this time the parent exits and the child becomes
        # an independent process
        try:
            pid = os.fork()
            if pid > 0:
                sys.exit(0)
        except OSError, e:
            sys.stderr.write('fork #2 failed: ({0}) {1}'.format(e.errno,
                                                                e.strerror))
            sys.exit(1)

        # Set the streams for the process

        if (logger.isEnabledFor(logging.DEBUG)):
            si = file('/tmp/spawner.debug.' + str(os.getpid()), 'a+')
            so = file('/tmp/spawner.debug.' + str(os.getpid()), 'a+')
            se = file('/tmp/spawner.debug.' + str(os.getpid()), 'a+')
        else:
            si = file(os.devnull, 'w')
            so = file(os.devnull, 'w')
            se = file(os.devnull, 'w')

        os.dup2(si.fileno(), sys.stdin.fileno())
        os.dup2(so.fileno(), sys.stdout.fileno())
        os.dup2(se.fileno(), sys.stderr.fileno())

        # Since we're in a different PID, if we passed a timingdb handle in,
        # it wouldn't be valid.  So, get a new one.
        try:
            timingdb = open_timing_database()
            timingcursor = timingdb.cursor()

        except:
            e = sys.exc_info()[0]
            logger.exception('Error creating cursor : {0}'.format(e))

        logger.debug('Timer number {0}'.format(str(timer)))

        sql = '''SELECT ip_range, nmap_options FROM timer_details WHERE
                 timer_id = ?'''
        timingcursor.execute(sql, (timer,))

        nmap_options = timingcursor.fetchone()

        fd, tmpfile = tempfile.mkstemp()

        cmdline = cfg['nmap_binary'] + ' -oX ' + tmpfile + ' ' + \
                  nmap_options[1] + ' ' + nmap_options[0]

        # Prior to 2.7.3, shlex wouldn't handle unicode.
        if (sys.version_info < (2, 7, 3)):
            cmdline = cmdline.encode('ascii')

        optionlist = shlex.split(cmdline)

        logger.debug('Will spawn with options {0}'.format(str(optionlist)))

        spawned = Popen(optionlist)
        logger.debug('Spawned - {0}'.format(str(spawned.pid)))

        sql = '''INSERT INTO spawned (timer_id, start_time, pid, overtime)
                 VALUES (?, ?, ?, 0)'''
        timingcursor.execute(sql, (timer, start_time, spawned.pid))

        timingdb.commit()

        spawned.wait()
        logger.debug('Timer {0} completed.'.format(str(timer)))

        nmap_filename = os.path.join(cfg['output_dir'], 'nmap_xml.' + str(timer)
                                     + '.' +  str(start_time) + '.xml.gpg')

        cmd = (cfg['gpg_binary'], '--batch', '--yes', '--output', nmap_filename,
               '--encrypt', '--recipient', cfg['gpg_key'], tmpfile)

        logger.debug('Encrypting {0} via GPG ({1})'.format(tmpfile, str(cmd)))

        Popen(cmd).wait()

        logger.debug('Deleting temporary file - {0}'.format(tmpfile))

        os.remove(tmpfile)

        sys.exit(0)

def open_timing_database():
    logger.info('Opening timing database')
    if (os.path.exists(cfg['datadir'])):
        try:
            timingdb = sqlite3.connect(os.path.join(cfg['datadir'],
                                                    cfg['timingdb']))

            timingcursor = timingdb.cursor()

            # Create the tables, if they don't exist already
            sql = '''CREATE TABLE IF NOT EXISTS timers
                                ( timer_id                  INT,
                                minute                      TEXT,
                                hour                        TEXT,
                                day                         TEXT,
                                month                       TEXT
                                )'''
            timingcursor.execute(sql)

            sql = '''CREATE TABLE IF NOT EXISTS timer_details
                                ( timer_id                  INT,
                                override_flag               BOOLEAN,
                                ip_range                    TEXT,
                                nmap_options                TEXT
                                )'''
            timingcursor.execute(sql)

            sql = '''CREATE TABLE IF NOT EXISTS spawned
                                ( timer_id                  INT,
                                start_time                  INT,
                                pid                         INT,
                                overtime                    BOOLEAN
                                )'''
            timingcursor.execute(sql)

            sql = '''CREATE TABLE IF NOT EXISTS spawn_log
                                ( timer_id                  INT,
                                start_time                  INT,
                                end_time                    INT,
                                status                      INT
                                )'''
            timingcursor.execute(sql)

        except:
            e = sys.exc_info()[0]
            logger.exception('Error accessing/creating sqlite database : {0}'
                             .format(e))
            sys.exit()

        return(timingdb)
    else:
        logger.critical('Data Directory does not exist')
        sys.exit()

###
# Usage
###
def usage():
    print ('Usage: ' + sys.argv[0] + ' [OPTION]... ');
    print '{0}'.format(programname)
    print
    print ('Mandatory arguments to long options are mandatory for short ' \
           'options too.')
    print ('  -c  <file>     configuration file')
    print ('  -d, --debug    debug')
    print ('  -h, --help     display this help and exit')
    print ('  -v, --verbose  verbose output')
    print ('  -V, --version  output version information and exit')
    print ('      --license  output license information and exit')

###
# License
###
def license():
    print('Copyright (C) 2014  Jason Frisvold <friz@godshell.com>');
    print('License GPLv2+: GNU GPL version 2 or later ' \
          '<http://gnu.org/licenses/gpl.html>.');
    print('This is free software: you are free to change and redistribute it.');
    print('There is NO WARRANTY, to the extent permitted by law.');
    print('');
    print('Written by Jason Frisvold.');

###
# Run the main program, pass in the CLI arguments (minus the command name)
try:
    if __name__ == "__main__":
        main(sys.argv[1:])
except KeyboardInterrupt:
    sys.exit()