#!/usr/bin/env python3
"""
poc_infer_gaps.py : tentative de combler, PAR INFERENCE HEURISTIQUE (pas
par cryptanalyse), le tiers d'octets que poc_recover_plaintext.py ne peut
pas recuperer.

IMPORTANT, a ne pas perdre de vue : le tiers d'octets manquant est
reellement chiffre par XChaCha20 (voir writeup.md section 3). Rien dans ce
script ne casse ce chiffre : XChaCha20 est un flux pseudo-aleatoire dont
chaque octet de keystream est independant des autres (c'est precisement
ce qui en fait un bon chiffre de flux). Connaitre le clair aux positions
voisines ne donne AUCUNE information mathematique sur le keystream a la
position manquante. Il n'existe donc pas de "correction" cryptographique
possible ici, seulement de la devination.

Ce que fait reellement ce script : pour du contenu textuel (langue
naturelle, code source), la structure du motif de Prince laisse TOUJOURS
les deux voisins immediats (gauche et droit) de chaque octet manquant en
clair (motif ENC,RAW,RAW,ENC,RAW,RAW,...). Pour un mot naturel, connaitre
tous les caracteres sauf un suffit tres souvent a le determiner par
recherche dans un dictionnaire (comme un mot croise a une seule case
vide). C'est une INFERENCE PAR CONTEXTE, pas un dechiffrement : elle
echoue silencieusement (et le script le signale) sur du contenu a haute
entropie (donnees binaires, identifiants, nombres, texte deja compresse
ou chiffre), qui reste et restera irrecuperable sans la cle privee de la
campagne.

Necessite le binaire `hunspell` et des dictionnaires locaux (ici : ceux
fournis avec OnlyOffice, fr_FR + en_US). Sans dictionnaire adapte a la
langue et au vocabulaire reel du fichier (jargon technique, noms propres),
le taux de succes chute nettement : c'est attendu, pas un bug de ce
script.

Usage : poc_infer_gaps.py FICHIER_RECUPERE_PARTIEL FICHIER_SORTIE
"""
import string
import subprocess
import sys

DICT_PATHS = [
    "/opt/onlyoffice/desktopeditors/dictionaries/fr_FR/fr_FR",
    "/opt/onlyoffice/desktopeditors/dictionaries/en_US/en_US",
]
ALPHABET = list(string.ascii_lowercase) + list("éèêëàâäùûüôöîïç")
GAP_BYTE = 0x00


def hunspell_check(words, dict_path):
    if not words:
        return {}
    proc = subprocess.run(
        ["hunspell", "-d", dict_path, "-a"],
        input="\n".join(words) + "\n",
        capture_output=True, text=True, encoding="utf-8", errors="replace",
    )
    lines = [l for l in proc.stdout.splitlines() if l and not l.startswith("@")]
    results = {}
    for word, line in zip(words, lines):
        results[word] = line.strip() == "*"
    return results


def is_word_char(b):
    if 65 <= b <= 90 or 97 <= b <= 122:
        return True
    try:
        c = bytes([b]).decode("latin-1")
    except Exception:
        return False
    return c.isalpha()


def main():
    if len(sys.argv) != 3:
        sys.exit(__doc__)
    src, dst = sys.argv[1], sys.argv[2]

    with open(src, "rb") as f:
        data = bytearray(f.read())

    # decoupage en "mots" (suites de caracteres alphabetiques, latin-1)
    words = []  # (start, end, has_gap, gap_pos)
    i = 0
    n = len(data)
    while i < n:
        if is_word_char(data[i]) or data[i] == GAP_BYTE:
            j = i
            gap_positions = []
            while j < n and (is_word_char(data[j]) or data[j] == GAP_BYTE):
                if data[j] == GAP_BYTE:
                    gap_positions.append(j)
                j += 1
            if len(gap_positions) == 1 and (j - i) >= 2:
                words.append((i, j, gap_positions[0]))
            i = j
        else:
            i += 1

    print(f"{len(words)} mot(s) avec exactement une lacune, longueur >= 2, candidats a l'inference")

    # generation de tous les candidats pour tous les mots, verification en un seul appel hunspell par dictionnaire
    candidates_by_word = {}
    all_candidates = set()
    for (start, end, gap) in words:
        cands = []
        for letter in ALPHABET:
            w = bytes(data[start:gap]) + letter.encode("latin-1") + bytes(data[gap + 1:end])
            try:
                w_str = w.decode("latin-1")
            except Exception:
                continue
            cands.append((letter, w_str))
            all_candidates.add(w_str)
        candidates_by_word[(start, end, gap)] = cands

    all_candidates = sorted(all_candidates)
    valid = {}
    for dict_path in DICT_PATHS:
        res = hunspell_check(all_candidates, dict_path)
        for w, ok in res.items():
            if ok:
                valid[w] = True

    n_inferred = 0
    n_ambiguous = 0
    n_no_match = 0
    for (start, end, gap), cands in candidates_by_word.items():
        matches = [letter for letter, w_str in cands if valid.get(w_str)]
        if len(matches) == 1:
            data[gap] = matches[0].encode("latin-1")[0]
            n_inferred += 1
        elif len(matches) == 0:
            n_no_match += 1
        else:
            n_ambiguous += 1

    n_still_gap = data.count(GAP_BYTE)
    total = len(data)
    print(f"\n{n_inferred} lacune(s) comblee(s) par inference dictionnaire (match unique)")
    print(f"{n_ambiguous} lacune(s) ambigue(s) (plusieurs mots valides possibles, non comblee)")
    print(f"{n_no_match} lacune(s) sans mot valide trouve dans le dictionnaire (non comblee)")
    print(f"{n_still_gap} octet(s) au total encore inconnus sur {total} "
          f"({100.0 * (total - n_still_gap) / total:.1f}% du fichier reconstitue, "
          f"exact + infere)")

    with open(dst, "wb") as f:
        f.write(data)
    print(f"ecrit dans {dst}")


if __name__ == "__main__":
    main()