#!/usr/bin/env python3
import os

from koalak.subcommand_parser import SubcommandParser

from pentest_helper.consts import (
    DEFAULT_RATE,
    FILE_USERS,
    FILE_PASSWORDS,
    FILE_WORDLIST,
    NAME_NMAP_TCP_FULL_SIMPLE,
    NAME_NMAP_UDP_TOP_SIMPLE,
    NAME_NMAP_UDP_OPEN_SC_SV,
    NAME_NMAP_TCP_OPEN_SC_SV,
)

from pentest_helper.init_functions import init_mission, init_host

from pentest_helper.consolidate_functions import consolidate_all

from pentest_helper.nmap_functions import (
    get_nmap_host_path,
    get_nmap_host_path_name,
    get_ports_from_gnmap,
)

from pentest_helper.utils import (
    check_mission_structure_initiated,
    run_cmd_then_print,
    check_running_with_root,
    logger,
)

# =========== #
# CMD PARSERS #
# =========== #
cmd_pentest_helper = SubcommandParser()

# Adding commands groups #
# ====================== #
group_utils = cmd_pentest_helper.add_group("utils")
group_cme = cmd_pentest_helper.add_group(
    "cme", description="Crackmapexec commands related"
)
group_nmap = cmd_pentest_helper.add_group("nmap", description="Nmap commands related")

# ============= #
# UTILS PARSERS #
# ============= #

# INIT CMD #
# ======== #
cmd_init = cmd_pentest_helper.add_subcommand(
    "init", group=group_utils, description="Init a new mission folders structure"
)
cmd_init.add_argument(
    "mission_path", help="Path of the mission ex: 'BOX01' or '/home/kali/mission1'"
)
cmd_init.add_argument(
    "--hosts", help="List of hosts to init with the mission", nargs="+"
)


@cmd_init
def run_init_mission(args):
    mission_path = args.mission_path
    hosts = args.hosts
    init_mission(mission_path, hosts=hosts)


# INIT HOST CMD #
# ============= #
cmd_init_host = cmd_pentest_helper.add_subcommand(
    "init_host",
    group=group_utils,
    description="Init a new host folders structure inside of a mission folders structure",
)
cmd_init_host.add_argument("name", help="IP or Domain Name of the Host")

@cmd_init_host
def run_init_host(args):
    name = args.name
    current_path = os.getcwd()
    init_host(current_path, name)


# CONSOLIDATE CMD #
# =============== #
cmd_consolidate = cmd_pentest_helper.add_subcommand(
    "consolidate",
    group=group_utils,
    description="Consolidate and update all files to be coherent."
    "Merge all credentials files (users passwords) into wordlist.txt",
)
cmd_consolidate.add_argument(
    "--fast",
    action="store_true",
    help="Do not heavy consolidation (ex: copying rockyou to custom rockou",
)

@cmd_consolidate
def run_consolidate(args):
    check_mission_structure_initiated()
    path = os.getcwd()
    fast = args.fast
    consolidate_all(path, fast=fast)


# ============ #
# NMAP PARSERS #
# ============ #
def add_nmap_args(parser):
    parser.add_argument("target")
    parser.add_argument("--rate", default=DEFAULT_RATE)

# NMAP_HOST CMD #
# ============= #
cmd_nmap_host = cmd_pentest_helper.add_subcommand(
    "nmap_host",
    description="Scan a host with nmap: tcp full port, tcp sc sv on open port, udp 1000 ports",
    group=group_nmap,
)
add_nmap_args(cmd_nmap_host)

@cmd_nmap_host
def run_nmap_host(args):
    run_nmap_host_tcp_full(args)
    run_nmap_host_tcp_sc_sv(args)
    run_nmap_host_udp_1000(args)
    run_nmap_host_udp_sc_sv(args)


# TCP FULL SCAN #
# ============= #
cmd_nmap_host_full_tcp = cmd_pentest_helper.add_subcommand(
    "nmap_host_tcp_full", group=group_nmap
)
add_nmap_args(cmd_nmap_host_full_tcp)

# TCP SC SV SCAN
cmd_nmap_host_tcp_sc_sv = cmd_pentest_helper.add_subcommand(
    "nmap_host_tcp_sc_sv", group=group_nmap
)
add_nmap_args(cmd_nmap_host_tcp_sc_sv)


cmd_nmap_host_udp_1000 = cmd_pentest_helper.add_subcommand(
    "nmap_host_udp_1000", group=group_nmap
)
add_nmap_args(cmd_nmap_host_udp_1000)

# UDP SC SV
cmd_nmap_host_udp_sc_sv = cmd_pentest_helper.add_subcommand(
    "nmap_host_udp_sc_sv", group=group_nmap
)
add_nmap_args(cmd_nmap_host_udp_sc_sv)


# CME PARSERS #
# =========== #
cmd_cme_smb_brute_users_and_passwords = cmd_pentest_helper.add_subcommand(
    "cme_brute_users_passwords", group=group_cme
)
cmd_cme_smb_brute_users_and_passwords.add_argument("target")


cmd_cme_smb_brute_wordlist = cmd_pentest_helper.add_subcommand(
    "cme_brute_wordlist", group=group_cme
)
cmd_cme_smb_brute_wordlist.add_argument("target")


# NMAP PARSERS #
# ============ #



# FULL NMAP SCAN









@cmd_cme_smb_brute_users_and_passwords
def run_cme_smb_bruteforce_users_and_passwords(args):
    target = args.target
    # crackmapexec smb 192.168.229.248 -u users.txt -p passwords.txt --continue-on-success
    run_cmd_then_print(
        [
            "crackmapexec",
            "smb",
            target,
            "-u",
            FILE_USERS,
            "-p",
            FILE_PASSWORDS,
            "--continue-on-success",
        ]
    )


@cmd_cme_smb_brute_wordlist
def run_cme_smb_bruteforce_wordlist(args):
    target = args.target
    # crackmapexec smb 192.168.229.248 -u users.txt -p passwords.txt --continue-on-success
    run_cmd_then_print(
        [
            "crackmapexec",
            "smb",
            target,
            "-u",
            FILE_WORDLIST,
            "-p",
            FILE_WORDLIST,
            "--continue-on-success",
        ]
    )


@cmd_nmap_host_full_tcp
def run_nmap_host_tcp_full(args):
    target = args.target
    rate = args.rate

    check_running_with_root()
    dir_path = get_nmap_host_path(target)
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    output_path = get_nmap_host_path_name(target, NAME_NMAP_TCP_FULL_SIMPLE)

    run_cmd_then_print(
        [
            "nmap",
            "-v",
            "-Pn",
            "-p-",
            f"--min-rate={rate}",
            "-T4",
            "-O",
            target,
            "-oA",
            output_path,
        ]
    )


@cmd_nmap_host_udp_1000
def run_nmap_host_udp_1000(args):
    target = args.target
    rate = args.rate

    check_running_with_root()
    dir_path = get_nmap_host_path(target)
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    output_path = get_nmap_host_path_name(target, NAME_NMAP_UDP_TOP_SIMPLE)

    run_cmd_then_print(
        [
            "nmap",
            "--open",
            "-vvv",
            "-n",
            "-Pn",
            "-sU",
            f"--min-rate={rate}",
            target,
            "-oA",
            output_path,
        ]
    )


@cmd_nmap_host_udp_sc_sv
def run_nmap_host_udp_sc_sv(args):
    target = args.target
    rate = args.rate

    check_running_with_root()
    dir_path = get_nmap_host_path(target)
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    output_path_gnmap = (
        get_nmap_host_path_name(target, NAME_NMAP_UDP_TOP_SIMPLE) + ".gnmap"
    )
    ports = get_ports_from_gnmap(output_path_gnmap, udp=True)
    if not ports:
        logger.info("No ports identified for -sC -sV")
        return

    logger.info(f"Scanning {target} with {ports}")
    output_path = get_nmap_host_path_name(target, NAME_NMAP_UDP_OPEN_SC_SV)
    # sudo nmap -p$ports -n -v -Pn --script=safe -sV  $rhost -oA
    run_cmd_then_print(
        [
            "nmap",
            "-n",
            "-v",
            "-Pn",
            "-p",
            ports,
            f"--min-rate={rate}",
            "-sC",
            "-sV",
            target,
            "-oA",
            output_path,
        ]
    )





@cmd_nmap_host_tcp_sc_sv
def run_nmap_host_tcp_sc_sv(args):
    target = args.target
    rate = args.rate

    check_running_with_root()
    dir_path = get_nmap_host_path(target)
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)
    output_path_gnmap = (
        get_nmap_host_path_name(target, NAME_NMAP_TCP_FULL_SIMPLE) + ".gnmap"
    )
    ports = get_ports_from_gnmap(output_path_gnmap)
    if not ports:
        logger.info("No ports identified for -sC -sV")
        return

    logger.info(f"Scanning {target} with {ports}")
    output_path = get_nmap_host_path_name(target, NAME_NMAP_TCP_OPEN_SC_SV)
    # sudo nmap -p$ports -n -v -Pn --script=safe -sV  $rhost -oA
    run_cmd_then_print(
        [
            "nmap",
            "-n",
            "-v",
            "-Pn",
            "-p",
            ports,
            f"--min-rate={rate}",
            "-sC",
            "-sV",
            target,
            "-oA",
            output_path,
        ]
    )


# =================== #
#         MAIN        #
# =================== #


def main():
    args = cmd_pentest_helper.parse_args()
    if args.subcommand != "init" and args.subcommand is not None:
        check_mission_structure_initiated()
    cmd_pentest_helper.run()


if __name__ == "__main__":
    main()
