#!/usr/bin/env python

import sys
import json
import urllib.parse

def generate_template(base_url, features):

    variables = {
        "IP_ADDRESS": "util.get_address()", 
        "BASE_URL": f'"{base_url}" if "LOCAL" not in sys.argv else "http://127.0.0.1:1337"',
        "PROXIES": json.dumps({"http":"http://127.0.0.1:8080", "https":"http://127.0.0.1:8080"})
    }
    
    if "proxies" in features or "burp" in features:
        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"

    request_method = 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)
"""

    methods = [request_method]

    if "login" in features or "account" in features:
        variables["USERNAME"] = '"Blindhero"'
        variables["PASSWORD"] = '"test1234"'
        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 "register" in features or "account" in features:
        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
""")

    main = """
if __name__ == "__main__":
    pass
"""

    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 template.py, for more information, visit
# https://git.romanh.de/Roman/HackingScripts

import os
import io
import re
import sys
import json
import time
import base64
import requests
import subprocess
import urllib.parse
from bs4 import BeautifulSoup
from hackingscripts import util, rev_shell
from hackingscripts.fileserver import HttpFileServer
from hackingscripts.sqli import MySQLi, PostgreSQLi
from urllib3.exceptions import InsecureRequestWarning
requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning)

{variables}

"""

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

if __name__ == "__main__":

    if len(sys.argv) < 2:
        print("Usage: %s <URL> [features]" % sys.argv[0])
        exit()

    url = sys.argv[1]
    if "://" not in url:
        url = "http://" + url

    features = [] if len(sys.argv) < 3 else sys.argv[2].split(",")
    template = generate_template(url, features)
    print(template)