# -*- coding: utf-8 -*-
# Thanks to the awesome tools: 
# - Windows-Exploit-Suggester: https://github.com/GDSSecurity/Windows-Exploit-Suggester
# - linux-exploit-suggester: https://github.com/mzet-/linux-exploit-suggester

from pupylib.PupyModule import *
from modules.lib.utils.shell_exec import shell_exec
import os
import subprocess
import random
import string
import tempfile

__class_name__="Exploit_Suggester"
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__),".."))

@config(compat=["linux", 'windows'], category="exploit")
class Exploit_Suggester(PupyModule):
    """ exploit suggester """
    
    def init_argparse(self):
        self.arg_parser = PupyArgumentParser(prog="Exploit_Suggester", description=self.__doc__)
        self.arg_parser.add_argument("--update", dest='update', action='store_true', default=False, help="Update Windows database (Internet access required on pupy server host)")
        self.arg_parser.add_argument("--shell", dest='shell', default="/bin/bash", help="Linux shell to use (default: /bin/bash)")

    def run(self, args):
        if args.update:
            self.execute_python_file('--update')
            return

        if self.client.is_windows():
            xls_file = self.find_xls_database(os.path.join(ROOT, "external", "Windows-Exploit-Suggester"))
            if not xls_file:
                self.error('no windows database file found. Please connect the server to internet and launch the --update command')
                return 

            # write systeminfo to the tmp directory
            random_file = ''.join([random.choice(string.ascii_lowercase) for x in range(0,random.randint(6,12))])
            full_path = os.path.join(tempfile.gettempdir(), random_file)
            f = open(full_path, 'w')
            f.write(shell_exec(self.client, "systeminfo").encode('utf-8', 'ignore'))
            f.close()

            self.success("Running Windows-Exploit-Suggester python script locally")
            self.execute_python_file('--database %s --systeminfo %s' % (xls_file, full_path))
            os.remove(full_path)

        elif self.client.is_linux():
            self.success("Running linux-exploit-suggester sh script on the target with the {0} shell on the target...".format(args.shell))
            if self.client.conn.modules.os.path.isfile(args.shell) == False:
                self.error("{0} does not exist on the target's system!".format(args.shell))
                self.error("You have to choose a valid shell")
                return -1
            code = open(os.path.join(ROOT, "external", "linux-exploit-suggester", "linux-exploit-suggester.sh"), 'r').read()
            p = self.client.conn.modules.subprocess.Popen(code, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE, shell=True, executable=args.shell)
            output, err = p.communicate()
            print output
             
    def execute_python_file(self, args):
        # Back up the current dir
        current_dir = os.getcwd()

        directory = os.path.join(ROOT, "external", "Windows-Exploit-Suggester")
        python_script = os.path.join(directory, "windows-exploit-suggester.py")

        # Change current dir
        os.chdir(directory)
        try:
            cmd = 'python %s %s' % (python_script, args)
            p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE, shell=True)
            output, err = p.communicate()
            print output
        except:
            pass
        
        # Restore the current directory path
        os.chdir(current_dir)


    def find_xls_database(self, path):
        xls_files = []
        for file in os.listdir(path):
            if not os.path.isdir(file) and file.endswith('.xls'):
                xls_files.append(file)
        xls_files.sort()
        
        # return the most recent xls file found 
        if xls_files:
            return xls_files[-1]
        else:
            return None
