#!/usr/bin/env python3
"""
xinux-crypt.py – Krypto-Demo für den Unterricht
================================================
Verwendung:
  python3 xinux-crypt.py --gen-asym-keys
  python3 xinux-crypt.py --sign <Name>
  python3 xinux-crypt.py --verify <Name> <pub.pem> <signatur.bin>
  python3 xinux-crypt.py --dh-init
  python3 xinux-crypt.py --dh-respond <dh_params.json>
  python3 xinux-crypt.py --dh-finish <dh_response.json>
  python3 xinux-crypt.py --encrypt <datei> --key <secret.bin>
  python3 xinux-crypt.py --decrypt <datei.enc> <datei.sha256> --key <secret.bin>

Kein Netzwerk – alle Dateien im aktuellen Verzeichnis.
Übertragung immer manuell per wget.
"""

import sys
import os
import hashlib
import json
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric.dh import (
    generate_parameters, DHParameterNumbers, DHPublicNumbers
)
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

SEP  = "=" * 55
SEP2 = "-" * 55

def info(text):  print(f"  {text}")
def ok(text):    print(f"  ✔  {text}")
def head(text):
    print()
    print(SEP)
    print(f"  {text}")
    print(SEP)

def fehler(text):
    print(f"\n  ✗  FEHLER: {text}")
    sys.exit(1)

def datei_pruefen(pfad):
    if not os.path.exists(pfad):
        fehler(f"Datei nicht gefunden: {pfad}")

def get_key_arg(args):
    """Liest --key <datei> aus der Argumentliste, Standard: secret.bin"""
    if "--key" in args:
        idx = args.index("--key")
        if idx + 1 < len(args):
            return args[idx + 1]
        else:
            fehler("--key benötigt einen Dateinamen als Argument")
    return "secret.bin"

# ──────────────────────────────────────────────────────────────
def gen_asym_keys():
    head("Asymmetrisches Schlüsselpaar generieren (RSA-2048)")

    priv = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    pub  = priv.public_key()

    with open("priv.pem", "wb") as f:
        f.write(priv.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.TraditionalOpenSSL,
            serialization.NoEncryption()
        ))

    pub_pem = pub.public_bytes(
        serialization.Encoding.PEM,
        serialization.PublicFormat.SubjectPublicKeyInfo
    )
    with open("pub.pem", "wb") as f:
        f.write(pub_pem)

    ok("priv.pem  – Private Key (geheim halten!)")
    ok("pub.pem   – Public Key  (kann jeder haben)")
    print()
    info("Inhalt von pub.pem:")
    print(SEP2)
    print(pub_pem.decode())
    print(SEP2)
    info("→ Gegenseite holt: wget http://<deine-IP>:8000/pub.pem")

# ──────────────────────────────────────────────────────────────
def sign(name):
    head(f"Namen signieren: '{name}'")

    datei_pruefen("priv.pem")

    with open("priv.pem", "rb") as f:
        priv = serialization.load_pem_private_key(f.read(), password=None)

    signatur = priv.sign(name.encode(), padding.PKCS1v15(), hashes.SHA256())

    with open("signatur.bin", "wb") as f:
        f.write(signatur)

    ok(f"'{name}' mit Private Key signiert")
    ok(f"signatur.bin geschrieben ({len(signatur)} Bytes)")
    print()
    info(f"Signatur (hex): {signatur.hex()[:48]}...")
    print()
    info("→ Gegenseite holt: wget http://<deine-IP>:8000/signatur.bin")
    info("  Gegenseite führt aus: python3 xinux-crypt.py --verify " + name)

# ──────────────────────────────────────────────────────────────
def verify(name, pub_datei, sig_datei):
    head(f"Signatur prüfen – erwarteter Name: '{name}'")

    datei_pruefen(pub_datei)
    datei_pruefen(sig_datei)

    with open(pub_datei, "rb") as f:
        pub = serialization.load_pem_public_key(f.read())

    with open(sig_datei, "rb") as f:
        signatur = f.read()

    info(f"{pub_datei} geladen ✔")
    info(f"{sig_datei} geladen ✔")
    info(f"Prüfe ob Signatur zu '{name}' passt...")
    print()

    try:
        pub.verify(signatur, name.encode(), padding.PKCS1v15(), hashes.SHA256())
        print(SEP2)
        info(f"✔  AUTHENTIFIZIERT: '{name}'")
        info(f"   Die Signatur ist gültig.")
        info(f"   Nur wer den passenden Private Key besitzt")
        info(f"   konnte diese Signatur erzeugen.")
        print(SEP2)
    except Exception:
        print(SEP2)
        info(f"✗  FEHLGESCHLAGEN: Signatur ungültig!")
        info(f"   Falscher Name oder falsche pub.pem")
        print(SEP2)

# ──────────────────────────────────────────────────────────────
def dh_init():
    head("Diffie-Hellman – Schritt 1: Parameter erzeugen (Initiator)")

    info("Generiere DH-Parameter p und g ...")
    params    = generate_parameters(generator=2, key_size=512)
    priv_dh   = params.generate_private_key()
    pub_dh    = priv_dh.public_key()

    p          = params.parameter_numbers().p
    g          = params.parameter_numbers().g
    pub_dh_int = pub_dh.public_numbers().y

    with open("dh_priv.pem", "wb") as f:
        f.write(priv_dh.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.PKCS8,
            serialization.NoEncryption()
        ))

    with open("dh_params.json", "w") as f:
        json.dump({"p": p, "g": g, "pub": pub_dh_int}, f)

    ok("dh_priv.pem    – eigener DH Private Key (bleibt lokal!)")
    ok("dh_params.json – p, g und eigener DH Public Value")
    print()
    info(f"p (Primzahl):      {str(p)[:50]}...")
    info(f"g (Generator):     {g}")
    info(f"DH Public Value:   {str(pub_dh_int)[:50]}...")
    print()
    info("→ Gegenseite holt: wget http://<deine-IP>:8000/dh_params.json")
    info("  Gegenseite führt aus: python3 xinux-crypt.py --dh-respond dh_params.json")

# ──────────────────────────────────────────────────────────────
def dh_respond(params_datei):
    head("Diffie-Hellman – Schritt 2: Antworten (Responder)")

    datei_pruefen(params_datei)

    with open(params_datei, "r") as f:
        data = json.load(f)

    p         = data["p"]
    g         = data["g"]
    pub_a_int = data["pub"]

    info(f"{params_datei} geladen ✔")
    info(f"p:                 {str(p)[:50]}...")
    info(f"g:                 {g}")
    info(f"Gegenseite PubVal: {str(pub_a_int)[:50]}...")
    print()

    param_numbers = DHParameterNumbers(p, g)
    params        = param_numbers.parameters()
    priv_dh       = params.generate_private_key()
    pub_dh        = priv_dh.public_key()
    pub_dh_int    = pub_dh.public_numbers().y

    pub_a     = DHPublicNumbers(pub_a_int, param_numbers).public_key()
    geheimnis = priv_dh.exchange(pub_a)
    aes_key   = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None, info=b"xinux-demo"
    ).derive(geheimnis)

    with open("dh_response.json", "w") as f:
        json.dump({"pub": pub_dh_int}, f)

    with open("secret.bin", "wb") as f:
        f.write(aes_key)

    ok("dh_response.json – eigener DH Public Value für Gegenseite")
    ok("secret.bin       – gemeinsames Geheimnis (AES-Key)")
    print()
    print(SEP2)
    info(f"Gemeinsamer AES-Key: {aes_key.hex()}")
    print(SEP2)
    print()
    info("→ Gegenseite holt: wget http://<deine-IP>:8000/dh_response.json")
    info("  Gegenseite führt aus: python3 xinux-crypt.py --dh-finish dh_response.json")

# ──────────────────────────────────────────────────────────────
def dh_finish(response_datei):
    head("Diffie-Hellman – Schritt 3: Geheimnis berechnen (Initiator)")

    datei_pruefen("dh_priv.pem")
    datei_pruefen("dh_params.json")
    datei_pruefen(response_datei)

    with open("dh_params.json", "r") as f:
        params_data = json.load(f)

    with open(response_datei, "r") as f:
        response_data = json.load(f)

    with open("dh_priv.pem", "rb") as f:
        priv_dh = serialization.load_pem_private_key(f.read(), password=None)

    p         = params_data["p"]
    g         = params_data["g"]
    pub_b_int = response_data["pub"]

    info(f"{response_datei} geladen ✔")
    info(f"Gegenseite PubVal: {str(pub_b_int)[:50]}...")
    print()

    param_numbers = DHParameterNumbers(p, g)
    pub_b         = DHPublicNumbers(pub_b_int, param_numbers).public_key()
    geheimnis     = priv_dh.exchange(pub_b)
    aes_key       = HKDF(
        algorithm=hashes.SHA256(), length=32, salt=None, info=b"xinux-demo"
    ).derive(geheimnis)

    with open("secret.bin", "wb") as f:
        f.write(aes_key)

    ok("secret.bin – gemeinsames Geheimnis (AES-Key)")
    print()
    print(SEP2)
    info(f"Gemeinsamer AES-Key: {aes_key.hex()}")
    print(SEP2)
    info("→ Beide Seiten haben denselben Key – ohne ihn je übertragen zu haben!")

# ──────────────────────────────────────────────────────────────
def encrypt(dateiname, key_datei):
    head(f"Datei verschlüsseln: {dateiname}")

    datei_pruefen(key_datei)
    datei_pruefen(dateiname)

    info(f"Schlüsseldatei: {key_datei}")

    with open(key_datei, "rb") as f:
        aes_key = f.read()

    with open(dateiname, "rb") as f:
        klartext = f.read()

    # SHA256-Hash des Klartexts
    klartext_hash = hashlib.sha256(klartext).hexdigest()

    # AES-CBC verschlüsseln
    iv      = os.urandom(16)
    pad_len = 16 - (len(klartext) % 16)
    padded  = klartext + bytes([pad_len] * pad_len)
    cipher  = Cipher(algorithms.AES(aes_key), modes.CBC(iv))
    enc     = cipher.encryptor()
    ciphertext = enc.update(padded) + enc.finalize()

    ausgabe      = dateiname + ".enc"
    hash_ausgabe = dateiname + ".sha256"

    with open(ausgabe, "wb") as f:
        f.write(iv + ciphertext)

    with open(hash_ausgabe, "w") as f:
        f.write(klartext_hash + "  " + dateiname + "\n")

    info(f"Klartext:         '{klartext.decode(errors='replace')[:60]}'")
    info(f"Größe:            {len(klartext)} Bytes")
    print()
    info(f"SHA256-Hash:      {klartext_hash}")
    info(f"(identisch mit:   sha256sum {dateiname})")
    print()
    info(f"IV (hex):         {iv.hex()}")
    info(f"Ciphertext (hex): {ciphertext.hex()[:48]}...")
    print()
    ok(f"{ausgabe}       – IV + Ciphertext")
    ok(f"{hash_ausgabe}  – SHA256-Hash des Klartexts")
    print()
    info(f"→ Gegenseite holt beide Dateien per wget")
    info(f"  python3 xinux-crypt.py --decrypt {ausgabe} {hash_ausgabe} --key {key_datei}")

# ──────────────────────────────────────────────────────────────
def decrypt(dateiname, hash_datei, key_datei):
    head(f"Datei entschlüsseln: {dateiname}")

    datei_pruefen(key_datei)
    datei_pruefen(dateiname)
    datei_pruefen(hash_datei)

    info(f"Schlüsseldatei: {key_datei}")

    basis = dateiname.replace(".enc", "")

    with open(key_datei, "rb") as f:
        aes_key = f.read()

    with open(dateiname, "rb") as f:
        rohdaten = f.read()

    with open(hash_datei, "r") as f:
        erwarteter_hash = f.read().strip().split()[0]

    iv         = rohdaten[:16]
    ciphertext = rohdaten[16:]

    info(f"IV (hex):         {iv.hex()}")
    info(f"Ciphertext (hex): {ciphertext.hex()[:48]}...")
    print()

    cipher    = Cipher(algorithms.AES(aes_key), modes.CBC(iv))
    dec       = cipher.decryptor()
    decrypted = dec.update(ciphertext) + dec.finalize()

    pad_len   = decrypted[-1]
    decrypted = decrypted[:-pad_len]

    tatsaechlicher_hash = hashlib.sha256(decrypted).hexdigest()

    ausgabe = basis + ".dec"
    with open(ausgabe, "wb") as f:
        f.write(decrypted)

    info(f"Klartext:         '{decrypted.decode(errors='replace')[:60]}'")
    print()
    info(f"Erwarteter Hash:    {erwarteter_hash}")
    info(f"Tatsächlicher Hash: {tatsaechlicher_hash}")
    print()

    if tatsaechlicher_hash == erwarteter_hash:
        print(SEP2)
        info("✔  INTEGRITÄT OK: Datei wurde nicht verändert!")
        info(f"   (identisch mit: sha256sum {ausgabe})")
        print(SEP2)
    else:
        print(SEP2)
        info("✗  INTEGRITÄT VERLETZT: Hash stimmt nicht überein!")
        info("   Die Datei wurde möglicherweise manipuliert.")
        print(SEP2)

    ok(f"{ausgabe} – entschlüsselte Datei gespeichert")

# ──────────────────────────────────────────────────────────────
# Hauptprogramm
# ──────────────────────────────────────────────────────────────
def usage():
    print("""
xinux-crypt.py – Krypto-Demo

  --gen-asym-keys                                  RSA-Schlüsselpaar generieren
  --sign <Name>                                    Namen mit Private Key signieren
  --verify <Name> <pub.pem> <sig.bin>              Signatur prüfen
  --dh-init                                        DH starten: p, g, pub erzeugen
  --dh-respond <dh_params.json>                    DH antworten + secret berechnen
  --dh-finish  <dh_response.json>                  DH abschließen + secret berechnen
  --encrypt <datei> --key <secret.bin>             Datei mit AES verschlüsseln + SHA256
  --decrypt <datei.enc> <datei.sha256> --key <secret.bin>  Datei entschlüsseln + Hash prüfen

  --key ist optional, Standard: secret.bin
""")
    sys.exit(1)

if len(sys.argv) < 2:
    usage()

cmd  = sys.argv[1]
args = sys.argv[1:]

if   cmd == "--gen-asym-keys":                                          gen_asym_keys()
elif cmd == "--sign"       and len(sys.argv) > 2:                       sign(sys.argv[2])
elif cmd == "--verify"     and len(sys.argv) > 4:                       verify(sys.argv[2], sys.argv[3], sys.argv[4])
elif cmd == "--dh-init":                                                 dh_init()
elif cmd == "--dh-respond" and len(sys.argv) > 2:                       dh_respond(sys.argv[2])
elif cmd == "--dh-finish"  and len(sys.argv) > 2:                       dh_finish(sys.argv[2])
elif cmd == "--encrypt"    and len(sys.argv) > 2:                       encrypt(sys.argv[2], get_key_arg(args))
elif cmd == "--decrypt"    and len(sys.argv) > 3:                       decrypt(sys.argv[2], sys.argv[3], get_key_arg(args))
else:                                                                    usage()