
import dsz, dsz.lp
import sys, random, time
import ops, util.ip, ops.timehelper
import os.path
import helper, scanbase
from math import floor
from datetime import datetime
import scanengine2
import monitorengine
import exceptions

def main(arguments):
    scanbase.setup_db()
    failout = False
    scansweepHelper = helper.scansweepHelper([x.lower() for x in arguments])
    scansweepHelper.check_env()
    create_mode = False
    if (scansweepHelper.options.database is not None):
        database_op = scansweepHelper.options.database
        if (not (database_op == 'create')):
            if (scansweepHelper.options.session is not None):
                scansweepHelper.database_display(database_op)
            else:
                scansweepHelper.database_display(database_op)
            return
        else:
            create_mode = True
    if (scansweepHelper.options.update is not None):
        scansweepHelper.handleupdate()
        return
    scanbase.write_metadata(scansweepHelper.scansweep_env, scansweepHelper.session, scansweepHelper.scansweep_logfile, scansweepHelper.scansweep_results, scansweepHelper.verbose)
    if (scansweepHelper.options.exclude is not None):
        scansweepHelper.parseexcludes(scansweepHelper.options.exclude)
    if (scansweepHelper.session == scansweepHelper.scansweep_env):
        if ((scansweepHelper.options.monitor is None) and (scansweepHelper.options.type is None)):
            dsz.ui.Echo('You must specify a type.', dsz.ERROR)
            return 0
        if ((scansweepHelper.options.monitor is None) and os.path.exists(scansweepHelper.options.type[0])):
            if (scansweepHelper.options.target is not None):
                dsz.ui.Echo('You cannot use -target when specifying a queue file.', dsz.ERROR)
                return 0
            queuefile = scansweepHelper.options.type[0]
            if (not scansweepHelper.verifyqueue(queuefile)):
                failout = True
            else:
                queuelist = scansweepHelper.getqueuefromfile(queuefile)
                for item in queuelist:
                    scansweepHelper.addtoqueue(item[0], item[1], scansweepHelper.scansweep_env)
        elif (scansweepHelper.options.type is not None):
            job_type = scansweepHelper.options.type[0].lower()
            job = '|'.join(scansweepHelper.options.type)
            if (not scansweepHelper.verifyjob(job_type, scansweepHelper.options.type)):
                dsz.ui.Echo('Invalid -type options, please verify your parameters.', dsz.ERROR)
                return 0
            candidate_list = []
            if (scansweepHelper.options.target is not None):
                if (type(scansweepHelper.options.target) == type([])):
                    for target_flag in scansweepHelper.options.target:
                        candidate_list.extend(scansweepHelper.parsetarget(target_flag))
                else:
                    candidate_list = scansweepHelper.parsetarget(scansweepHelper.options.target)
            else:
                dsz.ui.Echo('You must provide some targets with your scan.', dsz.ERROR)
                return 0
            if ((len(candidate_list) > 255) and (not scansweepHelper.options.cidroverride)):
                dsz.ui.Echo('You cannot specify more then 255 targets without the -cidroverride option', dsz.ERROR)
                failout = True
            else:
                scansweepHelper.addlisttoqueue({job: candidate_list})
        if (scansweepHelper.monitor is not None):
            for monitortype in scansweepHelper.monitor:
                if (scansweepHelper.verifymonitor(monitortype) is False):
                    dsz.ui.Echo(('%s is an invalid monitor type' % monitortype))
                    failout = True
        if ((scanbase.num_jobs(scansweepHelper.session) > 255) and (not scansweepHelper.options.cidroverride)):
            dsz.ui.Echo('You cannot specify more then 255 targets without the -cidroverride option', dsz.ERROR)
            failout = True
        if (scansweepHelper.options.escalate is not None):
            rulelist = scansweepHelper.parseescalate(scansweepHelper.options.escalate)
            if (len(rulelist) == 0):
                dsz.ui.Echo('You specified -escalate, but had only invalid rules. Exiting.', dsz.ERROR)
                failout = True
            for rule in rulelist:
                scantype = rule[1].split('|')[0]
                current_rulelist = scanbase.get_escalate_rules(scansweepHelper.session)
                if (rule not in current_rulelist):
                    scanbase.write_escalate_rule(scansweepHelper.session, rule)
                if (not (scantype == 'alert')):
                    scanbase.set_jobtype(scansweepHelper.session, scantype)
    elif ((scansweepHelper.options.type is not None) or (scansweepHelper.options.target is not None)):
        dsz.ui.Echo('You cannot specify -target or -type when using -session.', dsz.WARNING)
        failout = True
    else:
        dsz.ui.Echo('You are joining another session, and so will use the already available job queue and escalate rules.', dsz.WARNING)
    if (not scansweepHelper.verifytime(scanbase.get_jobtypes(scansweepHelper.session))):
        failout = True
    if failout:
        return 0
    scansweepHelper.printconfig()
    if create_mode:
        dsz.ui.Echo('Ran in create mode. Exiting.', dsz.WARNING)
        return
    dsz.lp.RecordToolUse('scansweep', scansweepHelper.toolversion, usage='EXERCISED', comment=' '.join([x.lower() for x in arguments]))
    try:
        scan(scansweepHelper)
    finally:
        dsz.ui.Echo(('=' * 100))
        scansweepHelper.showstats()
        print '\n\n'
        scansweepHelper.generateresults(quiet=False)

def scan(scansweepHelper):
    lastresults = 0
    alreadyoutput = []
    num_remaining = scanbase.num_jobs(scansweepHelper.session)
    sanity_string = ('[%s] Sanity output: %s jobs remaining, %s-%s remaining' % (dsz.Timestamp(), num_remaining, ops.timehelper.get_age_from_seconds((num_remaining * scansweepHelper.min_seconds)), ops.timehelper.get_age_from_seconds((num_remaining * scansweepHelper.max_seconds))))
    dsz.ui.Echo(sanity_string, dsz.GOOD)
    scansweepHelper.showstats()
    if (not os.path.exists(os.path.dirname(scansweepHelper.scansweep_logfile))):
        os.mkdir(os.path.dirname(scansweepHelper.scansweep_logfile))
    with open(scansweepHelper.scansweep_logfile, 'a') as f:
        f.write(('%s\n' % sanity_string))
    delta = time.time()
    scantime = time.time()
    originaltime = time.time()
    if (scansweepHelper.monitor is not None):
        scansweepHelper.activatemonitors()
    while True:
        if ((time.time() - originaltime) > scansweepHelper.maxtime):
            dsz.ui.Echo(('Maxtime of %s has been exceeded. Exiting.' % ops.timehelper.get_age_from_seconds(scansweepHelper.maxtime)), dsz.ERROR)
            break
        scan_job = scanbase.get_job(scansweepHelper.session)
        if (scan_job == False):
            if (scansweepHelper.monitor is None):
                break
        else:
            try:
                target = scan_job[1]
                job_info = scan_job[0].split('|')
                job_type = job_info[0]
                if (not util.ip.validate(target)):
                    target = scansweepHelper.resolvehostname(target)
                if (target == None):
                    continue
                target_scanner = scanengine2.get_scanengine(scan_job, scansweepHelper.timeout)
                target_scanner.execute_scan(False)
                if target_scanner.multiple_responses:
                    multi_response = target_scanner.return_data()
                    for response in multi_response:
                        scanbase.write_result(scansweepHelper.session, response.scan_type, response.target, response.return_data(), response.success, scan_job[0])
                else:
                    scanbase.write_result(scansweepHelper.session, target_scanner.scan_type, target_scanner.target, target_scanner.return_data(), target_scanner.success, scan_job[0])
                if target_scanner.success:
                    succ_out_string = ('[%s] %s (%s jobs remaining)' % (target_scanner.timestamp, target_scanner.return_success_message(), scanbase.num_jobs(scansweepHelper.session)))
                    dsz.ui.Echo(succ_out_string)
                    with open(scansweepHelper.scansweep_logfile, 'a') as f:
                        f.write(('%s\n' % succ_out_string))
                rulelist = scanbase.get_escalate_rules(scansweepHelper.session)
                for rule in rulelist:
                    if target_scanner.check_escalation(rule[0]):
                        if (rule[1] == 'alert'):
                            if (target_scanner.success == True):
                                esc_output_string = ('[%s]\t\tAlerting on %s by rule: (%s->%s)' % (dsz.Timestamp(), target, rule[0], rule[1]))
                            else:
                                esc_output_string = ('[%s] Alerting on %s by rule: (%s->%s)' % (dsz.Timestamp(), target, rule[0], rule[1]))
                            scansweepHelper.alert(esc_output_string)