#!/usr/bin/env python2

__author__ = "Anthony Miller-Rhodes (@amillerrhodes)"

import uuid
import argparse
import sys
# prevent creation of compiled bytecode files
sys.dont_write_bytecode = True
from recon.core import base

def recon_rpc(args):
    if args.server_type.lower() == 'xmlrpc':
        from SimpleXMLRPCServer import SimpleXMLRPCServer
        RPCServer = SimpleXMLRPCServer
        server = RPCServer((args.address, args.port), allow_none=True)
    elif args.server_type.lower() == 'jsonrpc':
        from jsonrpclib.SimpleJSONRPCServer import SimpleJSONRPCServer
        RPCServer = SimpleJSONRPCServer
        server = RPCServer((args.address, args.port))
    else:
        print('[!] Invalid RPC server type \'%s\'.' % (args.server_type))
        return
    server.register_multicall_functions()
    server.register_instance(ReconState())
    print("[+] Serving on %s:%d" % (args.address, args.port))
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print('')

class ReconState:

    def __init__(self):
        self.sessions = {}

    def init(self):
        sid = str(uuid.uuid4())
        self.sessions[sid] = {
            "recon": base.Recon(base.Mode.CLI),
            "module": None
        }
        self.sessions[sid]["module"] = self.sessions[sid]["recon"]
        return sid

    def _fetch_results(self, sid):
        results = self.sessions[sid]["module"].rpc_cache[:]
        self.sessions[sid]["module"].rpc_cache = []
        return results

    def use(self, param, sid):
        mod = self.sessions[sid]["recon"].do_use(param)
        self.sessions[sid]["module"] = mod

    def global_set(self, var, param, sid):
        self.sessions[sid]["recon"].do_set(var + " " + param)

    def set(self, var, param, sid):
        self.sessions[sid]["module"].do_set(var + " " + param)

    def unset(self, var, sid):
        self.sessions[sid]["module"].do_unset(var)

    def run(self, sid):
        self.sessions[sid]["module"].do_run(None)
        return self._fetch_results(sid)

    def add(self, table, param, sid):
        self.sessions[sid]["module"].do_add(table + " " + param)
        return self._fetch_results(sid)

    def delete(self, table, param, sid):
        self.sessions[sid]["module"].do_delete(table + " " + param)

    def show(self, param, sid):
        if param in self.sessions[sid]["module"].get_tables():
            return self.sessions[sid]["module"].query('SELECT ROWID, * FROM %s ORDER BY 1' % (param))

    def workspace(self, param, sid):
        self.sessions[sid]["recon"].init_workspace(param)

parser = argparse.ArgumentParser()
parser.add_argument("-t", type=str, action="store", default='jsonrpc', help="Set RPC server type", dest="server_type", metavar="[jsonrpc|xmlrpc]")
parser.add_argument("-a", type=str, action="store", default='0.0.0.0', help="Set RPC server bind address", dest="address", metavar="address")
parser.add_argument("-p", type=int, action="store", default=4141, help="Set RPC server port", dest="port", metavar="port")
args = parser.parse_args()
recon_rpc(args)
