import os
import sys
import shlex
import pickle
import socket
import inspect
import pkgutil
import subprocess
from pathlib import Path
from collections import defaultdict

os.environ["PYTHONPATH"] = f"{os.environ.get('PYTHONPATH')}:{str(Path(__file__).parent.resolve())}"

import cmd2
from cmd2.ansi import style

import recon


def get_scans():
    """ Iterates over the recon package and its modules to find all of the *Scan classes.

    * A contract exists here that says any scans need to end with the word scan in order to be found by this function.

    Returns:
        dict() containing mapping of {modulename: classname} for all potential recon-pipeline commands
    """
    scans = defaultdict(list)

    # recursively walk packages; import each module in each package
    for loader, module_name, is_pkg in pkgutil.walk_packages(recon.__path__, prefix="recon."):
        _module = loader.find_module(module_name).load_module(module_name)
        globals()[module_name] = _module

    # walk all modules, grabbing classes that we've written and add them to the classlist set
    for name, obj in inspect.getmembers(sys.modules[__name__]):
        if inspect.ismodule(obj) and not name.startswith("_"):
            for subname, subobj in inspect.getmembers(obj):
                if inspect.isclass(subobj) and subname.lower().endswith("scan"):
                    scans[subname].append(name)

    return scans


# tool definitions for the auto-installer
tools = {
    "go": {"installed": False, "dependencies": None, "commands": ["apt-get install -y -q golang"]},
    "luigi": {
        "installed": False,
        "dependencies": ["pipenv", "luigi-service"],
        "commands": ["pipenv install luigi"],
    },
    "luigi-service": {
        "installed": False,
        "dependencies": None,
        "commands": [
            f"cp {str(Path(__file__).parent / 'luigid.service')} /lib/systemd/system/luigid.service",
            f"cp $(which luigi) /usr/local/bin/luigid",
            "systemctl daemon-reload",
            "systemctl start luigid.service",
            "systemctl enable luigid.service",
        ],
        "shell": True,
    },
    "pipenv": {
        "installed": False,
        "dependencies": None,
        "commands": ["apt-get install -y -q pipenv"],
    },
    "masscan": {
        "installed": False,
        "dependencies": None,
        "commands": [
            "git clone https://github.com/robertdavidgraham/masscan /tmp/masscan",
            "make -s -j -C /tmp/masscan",
            "mv /tmp/masscan/bin/masscan /usr/local/bin/masscan",
            "rm -rf /tmp/masscan",
        ],
    },
    "amass": {
        "installed": False,
        "dependencies": None,
        "commands": ["apt-get install -y -q amass"],
    },
    "aquatone": {
        "installed": False,
        "dependencies": None,
        "commands": [
            "mkdir /tmp/aquatone",
            "wget -q https://github.com/michenriksen/aquatone/releases/download/v1.7.0/aquatone_linux_amd64_1.7.0.zip -O /tmp/aquatone/aquatone.zip",
            "unzip /tmp/aquatone/aquatone.zip -d /tmp/aquatone",
            "mv /tmp/aquatone/aquatone /usr/local/bin/aquatone",
            "rm -rf /tmp/aquatone",
        ],
    },
    "corscanner": {
        "installed": False,
        "dependencies": None,
        "commands": [
            "git clone https://github.com/chenjj/CORScanner.git /opt/CORScanner",
            "pip install -r /opt/CORScanner/requirements.txt",
            "pip install future",
        ],
    },
    "gobuster": {
        "installed": False,
        "dependencies": ["go"],
        "commands": [
            "go get github.com/OJ/gobuster",
            "(cd ~/go/src/github.com/OJ/gobuster && go build && go install)",
        ],
        "shell": True,
    },
    "tko-subs": {
        "installed": False,
        "dependencies": ["go"],
        "commands": [
            "go get github.com/anshumanbh/tko-subs",
            "(cd ~/go/src/github.com/anshumanbh/tko-subs && go build && go install)",
        ],
        "shell": True,
    },
    "subjack": {
        "installed": False,
        "dependencies": ["go"],
        "commands": [
            "go get github.com/haccer/subjack",
            "(cd ~/go/src/github.com/haccer/subjack && go build && go install)",
        ],
        "shell": True,
    },
    "webanalyze": {
        "installed": False,
        "dependencies": ["go"],
        "commands": [
            "go get github.com/rverton/webanalyze",
            "(cd ~/go/src/github.com/rverton/webanalyze && go build && go install)",
        ],
        "shell": True,
    },
    "recursive-gobuster": {
        "installed": False,
        "dependencies": None,
        "commands": [
            "git clone https://github.com/epi052/recursive-gobuster.git /opt/recursive-gobuster",
            "ln -s /opt/recursive-gobuster/recursive-gobuster.pyz /usr/local/bin",
        ],
    },
}

# options for ReconShell's 'scan' command
scan_parser = cmd2.Cmd2ArgumentParser()
scan_parser.add_argument("scantype", choices_function=get_scans)
scan_parser.add_argument(
    "--target-file",
    completer_method=cmd2.Cmd.path_complete,
    help="file created by the user that defines the target's scope; list of ips/domains",
)
scan_parser.add_argument(
    "--exempt-list", completer_method=cmd2.Cmd.path_complete, help="list of blacklisted ips/domains"
)
scan_parser.add_argument(
    "--wordlist", completer_method=cmd2.Cmd.path_complete, help="path to wordlist used by gobuster"
)
scan_parser.add_argument(
    "--interface",
    choices_function=lambda: [x[1] for x in socket.if_nameindex()],
    help="which interface masscan should use",
)
scan_parser.add_argument(
    "--recursive", action="store_true", help="whether or not to recursively gobust"
)
scan_parser.add_argument("--rate", help="rate at which masscan should scan")
scan_parser.add_argument(
    "--top-ports",
    help="ports to scan as specified by nmap's list of top-ports (only meaningful to around 5000)",
)
scan_parser.add_argument(
    "--ports", help="port specification for masscan (all ports example: 1-65535,U:1-65535)"
)
scan_parser.add_argument(
    "--threads", help="number of threads for all of the threaded applications to use"
)
scan_parser.add_argument("--scan-timeout", help="scan timeout for aquatone")
scan_parser.add_argument("--proxy", help="proxy for gobuster if desired (ex. 127.0.0.1:8080)")
scan_parser.add_argument("--extensions", help="list of extensions for gobuster (ex. asp,html,aspx)")
scan_parser.add_argument(
    "--local-scheduler",
    action="store_true",
    help="use the local scheduler instead of the central scheduler (luigid)",
)
scan_parser.add_argument(
    "--verbose",
    action="store_true",
    help="shows debug messages from luigi, useful for troubleshooting",
)

# options for ReconShell's 'install' command
install_parser = cmd2.Cmd2ArgumentParser()
install_parser.add_argument(
    "tool", help="which tool to install", choices=list(tools.keys()) + ["all"]
)


class ReconShell(cmd2.Cmd):
    prompt = "recon-pipeline> "

    @cmd2.with_argparser(scan_parser)
    def do_scan(self, args):
        """ Scan something. 
        
        Possible scans include 
            AmassScan           CORScannerScan      GobusterScan        SearchsploitScan
            ThreadedNmapScan    WebanalyzeScan      AquatoneScan        FullScan
            MasscanScan         SubjackScan         TKOSubsScan
        """
        scans = get_scans()
        command = ["luigi", "--module", scans.get(args.scantype)[0]]

        command.extend(args.__statement__.arg_list)

        sentry = False
        if args.verbose:
            command.pop(
                command.index("--verbose")
            )  # verbose is not a luigi option, need to remove it
            subprocess.run(command)
        else:
            # suppress luigi messages in favor of less verbose/cleaner output
            proc = subprocess.Popen(command, stderr=subprocess.PIPE, stdout=subprocess.PIPE)

            while True:
                output = proc.stderr.readline()
                if not output and proc.poll() is not None:
                    break
                if not output:
                    continue

                output = output.decode()
                if "===== Luigi Execution Summary =====" in output:
                    self.poutput()
                    sentry = True

                # block below used for printing status messages
                if sentry:
                    self.poutput(style(output.strip(), fg="bright_blue"))
                elif output.startswith("INFO") and output.split()[-1] == "PENDING":
                    words = output.split()
                    self.poutput(style(f"[-] {words[5].split('_')[0]} queued", fg="bright_white"))
                elif output.startswith("INFO") and "running" in output:
                    words = output.split()
                    # output looks similar to , pid=3938074) running   MasscanScan(
                    # want to grab the index of the luigi task running
                    scantypeidx = words.index("running") + 1
                    scantype = words[scantypeidx].split("(", 1)[0]
                    self.poutput(style(f"[*] {scantype} running...", fg="bright_yellow"))
                elif output.startswith("INFO") and output.split()[-1] == "DONE":
                    words = output.split()
                    self.poutput(
                        style(f"[+] {words[5].split('_')[0]} complete!", fg="bright_green")
                    )
        self.async_alert(
            style(
                "If anything went wrong, rerun your command with --verbose to enable debug statements.",
                fg="cyan",
                dim=True,
            )
        )

    @cmd2.with_argparser(install_parser)
    def do_install(self, args):
        """ Install any/all of the libraries/tools necessary to make the recon-pipeline function. """
        global tools

        cachedir = Path.home() / ".cache/"

        cachedir.mkdir(parents=True, exist_ok=True)

        persistent_tool_dict = cachedir / ".tool-dict.pkl"

        if args.tool == "all":
            for tool in tools.keys():
                self.do_install(tool)
            return

        if persistent_tool_dict.exists():
            tools = pickle.loads(persistent_tool_dict.read_bytes())

        if tools.get(args.tool).get("dependencies"):
            for dependency in tools.get(args.tool).get("dependencies"):
                if tools.get(dependency).get("installed"):
                    continue

                self.poutput(
                    style(
                        f"{args.tool} has an unmet dependency; installing {dependency}",
                        fg="blue",
                        bold=True,
                    )
                )
                self.do_install(dependency)

        if tools.get(args.tool).get("installed"):
            return self.poutput(style(f"{args.tool} is already installed.", fg="yellow"))
        else:
            self.poutput(style(f"Installing {args.tool}...", fg="blue", bold=True))

            for command in tools.get(args.tool).get("commands"):
                if tools.get(args.tool).get("shell"):  # go installs use subshells (...)
                    subprocess.run(command, shell=True)
                else:
                    subprocess.run(shlex.split(command))
            tools[args.tool]["installed"] = True

        self.poutput(style(f"{args.tool} installed!", fg="green", bold=True))
        pickle.dump(tools, persistent_tool_dict.open("wb"))


if __name__ == "__main__":
    rs = ReconShell(persistent_history_file="~/.reconshell_history")
    sys.exit(rs.cmdloop())