import sys
import os
from optparse import *
from coli import *

    
class BaseExploit(CommandlineWrapper):
    pocinfo = {
        'ID': '',
        'Name': 'The Poc Base Framework',
        'Author': 'w3h',
        'Created': '2015-01-08',
        'Description': '',

        'Vendor': '',
        'Device': '',
        'App': '',
        'Protocol': '',
        'References': {
            'CVE': '',
            'OSVDB': '',
            'CNVD': '',
            'CNNVD': ''
        },

        'Risk': '',  # H/M/L
        'VulType': ''
    }

    default_target = '127.0.0.1'
    default_port = None
    register_options = []
    base_options = [
        make_option('--verbose', action='store_true', dest='verbose',
                    default=False, help='Print verbose information.')
    ]

    def __init__(self):
        self.newframer = True

    def init_user_options(self):
        portflag = False
        targetflag = False
        for i in self.register_options:
            if i.get_opt_string()[2:] == 'TargetPort':
                portflag = True
            if i.get_opt_string()[2:] == 'TargetIp':
                targetflag = True

        if self.default_port and not portflag:
            newopt = make_option('-p', '--TargetPort', action='store', dest='TargetPort',
                                 type='int', default=int(self.default_port),
                                 help='The port of this poc [default:%d].' % int(self.default_port))
            self.base_options.insert(0, newopt)


        if self.default_target and not targetflag :
            newopt = make_option('-t', '--TargetIp', action='store', dest='TargetIp',
                                 type='string', default=None,
                                 help='The target of this poc.')
            self.base_options.insert(0, newopt)

        self.base_options += self.register_options

    def get_options(self):
        super(BaseExploit, self).get_options()
        self.init_user_options()
        self.get_coli_parser().add_options(self.base_options)

    def getParam(self, para):
        for (k, v) in self.options.__dict__.items():
            if not k == para:
                continue
            for i in self.base_options:
                if k == i.get_opt_string()[2:]:
                    if 'int' == i.type:
                        return int(v)
                    else:
                        return v

        return None

    def __shellCall(self, opts, args):
        context = {}
        # get user para and global Paramaters
        self.options = opts
        self.TargetIp = self.getParam('TargetIp')
        self.TargetPort = self.getParam('TargetPort')

        (fhNo, logConfig) = self.processWrapperParams(opts)
        retval = self.processParams(opts.__dict__, None, None, context, logConfig)

    def run(self, argv):
        """Effectively "main" from Commandlinewrapper"""
        self.get_options()
        try:
            (opts, args) = self.get_coli_parser().parse_args(argv)
            self.options = opts
            if not argv:
                self.get_coli_parser().print_help()
                return

            if opts.InConfig is None:
                # not xml file run mode
                self.__shellCall(opts, args)
                return

            super(BaseExploit, self).run(argv)
        except Exception as e:
            print "Failed: {0}".format(e)
            raise
   
    #
    # These need to be implemented by the exploit
    # 
    def processParams(self, inputs, constants, outputs, context, logConfig):
        """Process the input parameters and achieve the intended purpose"""
        try:
            # script pre check
            if not self.check():
                return

            return self.exploit(inputs, constants, outputs, context, logConfig)
        except Exception,e:
            raise

    def getID(self):
        """Return the plugin ID"""
        return 0

    def cleanup(self, flags, context, logConfig):
        """Cleanup any errant connections or data after the rendezvous is done"""
        return

    def validateParams(self, inputs):
        """Validate parameters to verify sane values"""
        return True

    def exploit(self, *args, **kwargs):
        """Process the input parameters and achieve the intended purpose"""
        raise NotImplementedError("processParams must be implemented")

    def check(self, *args, **kwargs):
        return True


def MainEntry(obj, flag):
    if flag == "__main__":
        obj().run(sys.argv[1:])
