#!/usr/bin/env python3
"""
سیمرغ (Simorgh) — Shamir 5-of-7 COMBINER (player-side, offline).

Bring any 5 of the 7 khan flags you have captured; this reconstructs the FINAL
flag. This tool + simorgh.json are safe to publish: without 5 correct flags it
reveals nothing.

Usage:
    python3 combine.py simorgh.json  YEK{...} YEK{...} YEK{...} YEK{...} YEK{...}
    # or run with no flag args and paste them one per line (Ctrl-D to finish).

This is the reference (authoritative) combiner. The Android app ships an
equivalent Kotlin GF(p) recover; if they ever disagree, THIS output wins.
"""
import base64
import hashlib
import json
import sys

from sslib import shamir
from Crypto.Cipher import AES
from Crypto.Util.Padding import unpad


def combine(simorgh: dict, flags):
    by_id = {r["id"]: r for r in simorgh["records"]}
    shares, used = [], set()
    for flag in flags:
        flag = flag.strip()
        if not flag:
            continue
        rid = hashlib.sha256(b"i" + flag.encode()).hexdigest()[:16]
        rec = by_id.get(rid)
        if rec is None:
            print(f"[-] no share matches flag: {flag!r} (wrong/typo'd?)", file=sys.stderr)
            continue
        if rid in used:            # same flag twice -> one share
            continue
        used.add(rid)
        key = hashlib.sha256(b"k" + flag.encode()).digest()
        raw = base64.b64decode(rec["blob"])
        share_str = unpad(AES.new(key, AES.MODE_CBC, raw[:16]).decrypt(raw[16:]), 16).decode()
        shares.append(share_str)

    need = simorgh["required_shares"]
    if len(shares) < need:
        raise SystemExit(f"[-] need {need} distinct correct flags, got {len(shares)}.")

    data = shamir.from_base64({
        "required_shares": need,
        "prime_mod": simorgh["prime_mod"],
        "shares": shares[:need],
    })
    return shamir.recover_secret(data).decode("utf-8")


def main(argv):
    if len(argv) < 2:
        raise SystemExit("usage: combine.py simorgh.json [flag ...]")
    with open(argv[1]) as f:
        simorgh = json.load(f)
    flags = argv[2:]
    if not flags:
        print("paste your khan flags, one per line (Ctrl-D to finish):", file=sys.stderr)
        flags = [ln for ln in sys.stdin.read().splitlines()]
    final = combine(simorgh, flags)
    print(final)


if __name__ == "__main__":
    main(sys.argv)