#!/usr/bin/env python

import re
import sys
import json
import argparse
import urllib.parse

def generate_template(base_url, features):

    # we could all need that
    imports = [
        "os", "io", "re", "sys",
        "json", "time", "base64", "requests",
        "subprocess", "urllib.parse"
    ]

    partial_imports = {
        "bs4": ["BeautifulSoup"],
        "hackingscripts": ["util", "rev_shell"],
        "urllib3.exceptions": ["InsecureRequestWarning"]
    }

    main_code = []
    methods = []
    ip_address_arg = next(filter(lambda f: re.match(r"ip_address=(.*)", f), features), None)
    ip_address = "util.get_address()" if not ip_address_arg else "'" + ip_address_arg[1] + "'"

    variables = {
        "IP_ADDRESS": ip_address, 
        "BASE_URL": f'"{base_url}" if "LOCAL" not in sys.argv else "http://127.0.0.1:1337"'
    }
    
    proxy_arg = next(filter(lambda f: re.match(r"proxy=(.*)", f), features), None)
    if proxy_arg or "burp" in features:
        proxy_url = "http://127.0.0.1:8080" if not proxy_arg else proxy_arg[1]
        variables["PROXIES"] = json.dumps({"http": proxy_url, "https": proxy_url})
        proxy = """
    if \"proxies\" not in kwargs:
        kwargs[\"proxies\"] = PROXIES
"""
    else:
        proxy = ""

    if "vhost" in features or "subdomain" in features:
        url_parts = urllib.parse.urlparse(base_url)
        host_name = url_parts.netloc
        variables["HOST_NAME"] = f"'{host_name}' if \"LOCAL\" not in sys.argv else \"127.0.0.1:1337\""
        vhost_param = ", vhost=None"
        full_url = f"f'{url_parts.scheme}://{{vhost}}.{{HOST_NAME}}{{uri}}' if vhost else BASE_URL + uri"
    else:
        vhost_param = ""
        full_url = "BASE_URL + uri"
        
    methods.insert(0, f"""def request(method, uri{vhost_param}, **kwargs):
    if not uri.startswith("/") and uri != "":
        uri = "/" + uri

    client = requests
    if "session" in kwargs:
        client = kwargs["session"]
        del kwargs["session"]
    
    if "allow_redirects" not in kwargs:
        kwargs["allow_redirects"] = False
    
    if "verify" not in kwargs:
        kwargs["verify"] = False
    {proxy}
    url = {full_url}
    return client.request(method, url, **kwargs)
""")

    if "register" in features or "account" in features:
        main_code.append("""if not register(USERNAME, PASSWORD):
        exit(1)
""")
        variables["USERNAME"] = '"Blindhero"'
        variables["PASSWORD"] = '"test1234"'
        methods.append("""
def register(username, password):
    res = request("POST", "/register", data={"username": username, "password": password})
    if res.status_code != 200:
        print("[-] Error registering")
        exit()
    
    return True
""")

    if "login" in features or "account" in features:
        main_code.append("""session = login(USERNAME, PASSWORD)
    if not session:
        exit(1)
""")
        variables["USERNAME"] = '"username"'
        variables["PASSWORD"] = '"password"'
        methods.append("""
def login(username, password):
    session = requests.Session()
    res = request("POST", "/login", data={"username": username, "password": password}, session=session)
    if res.status_code != 200:
        print("[-] Error logging in")
        exit()
    
    return session
""")

    if "sqli" in features:
        partial_imports["hackingscripts.sqli"] = ["MySQLi", "PostgreSQLi", "BlindSQLi", "ReflectedSQLi"]
        methods.append("""
class ReflectedSQLiPoC(MySQLi, ReflectedSQLi):
    def __init__(self):
        # TODO: specify reflected columns with their types
        super().__init__([None, str, int])
    def reflected_sqli(self, columns: list, table=None, condition=None, offset=None, verbose=False):
        # TODO: build query and extract columns from response
        return None
""")
        methods.append("""
class BlindSQLiPoC(MySQLi, BlindSQLi):
    def blind_sqli(self, condition: str, verbose=False) -> bool:
        # TODO: build query and evaluate condition
        return False
""")

        main_code.append("""poc = ReflectedSQLiPoC()
    print(poc.get_current_user())  
""")

    if "http-server" in features or "file-server" in features:
        partial_imports["hackingscripts.fileserver"] = ["HttpFileServer"]
        main_code.append("""file_server = HttpFileServer("0.0.0.0", 3000)
    file_server.enableLogging()
    file_server.addRoute("/dynamic", on_request)
    file_server.addFile("/static", b"static-content")
    file_server.startBackground()
""")

        methods.append("""
def on_request(req):
    # TODO: auto generated method stub
    return 200, b"", { "X-Custom-Header": "1" }     
""")

    if len(main_code) == 0:
        main_code = ["pass"]

    main = f"""
if __name__ == "__main__":
    {'\n    '.join(main_code)}
"""

    imports = "\n".join(f"import {i}" for i in sorted(imports, key=len))
    imports += "\n" + "\n".join(sorted(list(f"from {p} import {', '.join(i)}" for p, i in partial_imports.items()), key=len))
    variables = "\n".join(f"{k} = {v}" for k, v in variables.items())
    header = f"""#!/usr/bin/env python

#
# THE BASE OF THIS FILE WAS AUTOMATICALLY GENERATED BY {' '.join(sys.argv)}
# For more information, visit: https://git.romanh.de/Roman/HackingScripts
# 

{imports}
requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning)

{variables}

"""

    return header + "".join(methods) + main

if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="Exploit Template for web attacks",
        formatter_class=argparse.RawTextHelpFormatter
    )

    available_features = [
        "ip_address=[...]: Local IP-Address for reverse connections",
        "burp|proxy=[...]: Tunnel traffic through a given proxy or Burp defaults",
        "subdomain|vhost: Allow to specify a subdomain for outgoing requests",
        "register|account: Generate an account registration method stub",
        "login|account: Generate an account login method stub",
        "sqli: Generate an template SQL-Injection class",
        "http-server|file-server: Generate code for starting an in-memory http server"
    ]

    parser.add_argument("url", type=str, help="Target URL")
    parser.add_argument(
        "-f",
        "--features",
        nargs="*",
        type=str,
        default=[],
        help="Optional list of features:\n- " + "\n- ".join(available_features)
    )

    args = parser.parse_args()

    url = args.url
    if "://" not in url:
        url = "http://" + url

    features = args.features
    template = generate_template(url, features)
    print(template)