#!/usr/bin/env python3
"""
schattenpost — One-Time-Pad-Verschlüsselung + LSB-Steganografie.

Verschlüsselt Text mit einem echten One-Time-Pad (XOR) und versteckt den
Ciphertext in den niederwertigsten Bits eines Trägerfotos. Das Ergebnis sieht
aus wie ein ganz normales Bild, trägt aber die Nachricht bit-genau in sich.

RICHTUNGSGETRENNTE SCHLÜSSEL
----------------------------
Jede Beziehung nutzt ZWEI unabhängige Schlüsselströme — einen pro Senderichtung.
Dadurch können sich zwei gleichzeitig sendende Parteien nie überschneiden
(kein „two-time-pad"). Jede Seite hält:

  send.key  — dein Sendeschlüssel. Trägt den Sende-Offset IM Header; er wandert
              bei jeder Nachricht vor. Der Partner hat dieselben Bytes als recv.key.
  recv.key  — dein Empfangsschlüssel. Wird nur GELESEN (Offset kommt aus dem Bild),
              nie verändert.

Der Sende-Offset lebt also im Schlüssel selbst. Legt man den Schlüssel auf einen
USB-Stick und hängt ihn nur ans gerade benutzte Gerät, gibt es physisch nur einen
einzigen Offset-Zähler — kein Auseinanderdriften über mehrere Geräte möglich.

WICHTIG
-------
  * Stego-Bild IMMER als Datei/Dokument verschicken, nie als „Foto" — Messenger
    komprimieren Fotos verlustbehaftet und zerstören die versteckten Bits.
  * send.key niemals doppelt/parallel auf zwei Geräten bespielen (siehe oben).

Subcommands:
  genkey   Erzeugt ein gerichtetes Schlüsselpaar für eine Beziehung.
  hide     Verschlüsselt + versteckt eine Nachricht (braucht send.key).
  reveal   Extrahiert + entschlüsselt eine Nachricht (braucht recv.key).

Abhängigkeit: Pillow (pip install Pillow)
"""

import argparse
import os
import struct
import sys
from datetime import datetime

from PIL import Image

# Unauffälliger Default-Dateiname (Screenshot-Stil passt zu PNG; echte
# Kamerafotos sind JPG). Identisch zum Web-Frontend.
CARRIER_PREFIX = "Screenshot"


def default_carrier_name() -> str:
    return datetime.now().strftime(f"{CARRIER_PREFIX}_%Y%m%d_%H%M%S.png")

# --- Schlüsseldatei ---------------------------------------------------------
KEY_MAGIC = b"SPK1"
KEY_VERSION = 1
ROLE_RECV = 0
ROLE_SEND = 1
# magic(4) ver(1) role(1) stream_id(8) send_offset(uint64 BE, 8) + 2 pad = 24
KEY_FMT = ">4sBB8sQ2x"
KEY_HEADER_LEN = struct.calcsize(KEY_FMT)
KEY_OFFSET_POS = 14  # Byteposition des send_offset-Feldes (für In-Place-Update)

# --- Stego-Container (im Bild) ---------------------------------------------
# v3 authentifiziert zusätzlich mit einem Poly1305-Einmal-MAC (Tag im Header).
STEGO_MAGIC = b"SPS3"
STEGO_VERSION = 3
# Vor-Tag-Header: magic(4) ver(1) stream_id(8) offset(uint64 BE,8) len(uint32 BE,4)=25
STEGO_PRE_FMT = ">4sB8sQI"
STEGO_PRE_LEN = struct.calcsize(STEGO_PRE_FMT)        # 25
STEGO_TAG_LEN = 16                                    # Poly1305-Tag
STEGO_HEADER_LEN = STEGO_PRE_LEN + STEGO_TAG_LEN      # 41
MAC_KEY_LEN = 32                                      # r(16)+s(16), aus dem Pad

# Alt-Format (nur noch beim Aufdecken akzeptiert, unauthentifiziert).
STEGO_MAGIC_V2 = b"SPS2"
STEGO_VERSION_V2 = 2


# --------------------------------------------------------------------------- #
# One-Time-Pad
# --------------------------------------------------------------------------- #
def otp_xor(data: bytes, material: bytes, offset: int) -> bytes:
    if offset + len(data) > len(material):
        raise ValueError(
            f"Schlüssel erschöpft: ab Offset {offset} brauche ich {len(data)} "
            f"Bytes, habe nur {len(material) - offset}."
        )
    return bytes(d ^ material[offset + i] for i, d in enumerate(data))


# --------------------------------------------------------------------------- #
# Poly1305 — informationstheoretisch sicherer Einmal-MAC (RFC 8439)
# --------------------------------------------------------------------------- #
# Der 32-Byte-Schlüssel (r|s) wird pro Nachricht frisch aus dem One-Time-Pad
# gezogen (direkt hinter dem Nachrichten-Keystream) und NIE wiederverwendet —
# genau die Bedingung, unter der Poly1305 fälschungssicher ist, ohne die
# informationstheoretische Sicherheit des OTP aufzugeben.
_P1305_P = (1 << 130) - 5
_P1305_CLAMP = 0x0ffffffc0ffffffc0ffffffc0fffffff


def poly1305_mac(msg: bytes, key: bytes) -> bytes:
    if len(key) != 32:
        raise ValueError("Poly1305-Schlüssel muss 32 Byte sein.")
    r = int.from_bytes(key[:16], "little") & _P1305_CLAMP
    s = int.from_bytes(key[16:32], "little")
    acc = 0
    for i in range(0, len(msg), 16):
        block = msg[i:i + 16]
        n = int.from_bytes(block, "little") | (1 << (8 * len(block)))
        acc = (acc + n) * r % _P1305_P
    return ((acc + s) % (1 << 128)).to_bytes(16, "little")


def ct_equal(a: bytes, b: bytes) -> bool:
    """Konstant-zeitiger Vergleich (kein Early-Exit-Timing-Leck)."""
    if len(a) != len(b):
        return False
    diff = 0
    for x, y in zip(a, b):
        diff |= x ^ y
    return diff == 0


# --------------------------------------------------------------------------- #
# Schlüsseldateien
# --------------------------------------------------------------------------- #
def read_key(path: str) -> dict:
    with open(path, "rb") as fh:
        raw = fh.read()
    if len(raw) < KEY_HEADER_LEN:
        raise ValueError(f"{path}: keine gültige Schlüsseldatei (zu kurz).")
    magic, ver, role, stream_id, send_offset = struct.unpack(
        KEY_FMT, raw[:KEY_HEADER_LEN]
    )
    if magic != KEY_MAGIC:
        raise ValueError(f"{path}: keine schattenpost-Schlüsseldatei.")
    if ver != KEY_VERSION:
        raise ValueError(f"{path}: unbekannte Schlüsselversion {ver}.")
    return {
        "role": role,
        "stream_id": stream_id,
        "send_offset": send_offset,
        "material": raw[KEY_HEADER_LEN:],
    }


def write_key(path: str, role: int, stream_id: bytes, send_offset: int,
              material: bytes) -> None:
    header = struct.pack(KEY_FMT, KEY_MAGIC, KEY_VERSION, role, stream_id,
                         send_offset)
    with open(path, "wb") as fh:
        fh.write(header + material)


def advance_send_offset(path: str, new_offset: int) -> None:
    """Schreibt nur das 8-Byte-Offset-Feld in-place (kein Rewrite des Materials)."""
    with open(path, "r+b") as fh:
        fh.seek(KEY_OFFSET_POS)
        fh.write(struct.pack(">Q", new_offset))


# --------------------------------------------------------------------------- #
# LSB-Steganografie
# --------------------------------------------------------------------------- #
def _bits_from_bytes(data: bytes):
    for byte in data:
        for i in range(7, -1, -1):
            yield (byte >> i) & 1


def embed(carrier_path: str, payload: bytes, out_path: str) -> None:
    img = Image.open(carrier_path).convert("RGB")
    width, height = img.size
    need_bits = len(payload) * 8
    if need_bits > width * height * 3:
        raise ValueError(
            f"Trägerbild zu klein: brauche {need_bits // 8} Bytes, Platz für "
            f"{width * height * 3 // 8}. Nimm ein größeres Foto."
        )
    px = img.load()
    bits = _bits_from_bytes(payload)
    done = False
    for y in range(height):
        if done:
            break
        for x in range(width):
            r, g, b = px[x, y]
            ch = [r, g, b]
            for c in range(3):
                bit = next(bits, None)
                if bit is None:
                    done = True
                    break
                ch[c] = (ch[c] & ~1) | bit
            px[x, y] = (ch[0], ch[1], ch[2])
            if done:
                break
    img.save(out_path, "PNG")


def _read_bits(px, width, height, n_bits):
    out = bytearray()
    acc = count = read = 0
    for y in range(height):
        for x in range(width):
            for c in px[x, y]:
                acc = (acc << 1) | (c & 1)
                count += 1
                read += 1
                if count == 8:
                    out.append(acc)
                    acc = count = 0
                if read >= n_bits:
                    return bytes(out)
    return bytes(out)


def extract(stego_path: str) -> dict:
    img = Image.open(stego_path).convert("RGB")
    width, height = img.size
    px = img.load()
    # Vor-Tag-Header ist bei v2 und v3 identisch aufgebaut (25 B) -> erst lesen.
    pre = _read_bits(px, width, height, STEGO_PRE_LEN * 8)
    magic, ver, stream_id, offset, length = struct.unpack(STEGO_PRE_FMT, pre)
    if magic == STEGO_MAGIC and ver == STEGO_VERSION:
        blob = _read_bits(px, width, height,
                          (STEGO_HEADER_LEN + length) * 8)
        return {
            "version": ver,
            "stream_id": stream_id,
            "offset": offset,
            "length": length,
            "tag": blob[STEGO_PRE_LEN:STEGO_HEADER_LEN],
            "mac_input_header": blob[:STEGO_PRE_LEN],
            "ciphertext": blob[STEGO_HEADER_LEN:STEGO_HEADER_LEN + length],
        }
    if magic == STEGO_MAGIC_V2 and ver == STEGO_VERSION_V2:
        blob = _read_bits(px, width, height,
                          (STEGO_PRE_LEN + length) * 8)
        return {
            "version": ver,
            "stream_id": stream_id,
            "offset": offset,
            "length": length,
            "tag": None,
            "ciphertext": blob[STEGO_PRE_LEN:STEGO_PRE_LEN + length],
        }
    raise ValueError(
        "Kein schattenpost-Container gefunden (Magic/Version passt nicht). "
        "Wurde das Bild als Foto verlustbehaftet komprimiert?"
    )


# --------------------------------------------------------------------------- #
# Commands
# --------------------------------------------------------------------------- #
def cmd_genkey(args):
    size = args.size
    mine = args.output + "_meine_seite"
    theirs = args.output + "_partner_seite"
    for d in (mine, theirs):
        if os.path.exists(d) and not args.force:
            sys.exit(f"[!] {d} existiert schon. --force zum Überschreiben.")
        os.makedirs(d, exist_ok=True)

    # Zwei unabhängige Ströme: sid_out = ich->Partner, sid_in = Partner->ich
    sid_out, sid_in = os.urandom(8), os.urandom(8)
    mat_out, mat_in = os.urandom(size), os.urandom(size)

    # Meine Seite: sende mit Strom OUT, empfange Strom IN
    write_key(os.path.join(mine, "send.key"), ROLE_SEND, sid_out, 0, mat_out)
    write_key(os.path.join(mine, "recv.key"), ROLE_RECV, sid_in, 0, mat_in)
    # Partnerseite (spiegelverkehrt): sendet mit Strom IN, empfängt Strom OUT
    write_key(os.path.join(theirs, "send.key"), ROLE_SEND, sid_in, 0, mat_in)
    write_key(os.path.join(theirs, "recv.key"), ROLE_RECV, sid_out, 0, mat_out)

    print(f"[+] Schlüsselpaar erzeugt ({size} Bytes je Richtung, echter Zufall).")
    print(f"    {mine}/    <- DU behältst diesen Ordner")
    print(f"    {theirs}/  <- gib DIESEN Ordner dem Partner (Stick, persönlich)")
    print("[i] Jede Seite hat send.key (zum Senden) + recv.key (zum Empfangen).")
    print("[i] send.key nur an EINEM Gerät gleichzeitig benutzen.")


def cmd_hide(args):
    if args.text is not None:
        plaintext = args.text.encode("utf-8")
    elif args.infile:
        with open(args.infile, "rb") as fh:
            plaintext = fh.read()
    else:
        plaintext = sys.stdin.buffer.read()
    if not plaintext:
        sys.exit("[!] Leere Nachricht.")

    out_path = args.output or default_carrier_name()
    key = read_key(args.key)
    if key["role"] != ROLE_SEND:
        sys.exit("[!] Das ist ein Empfangsschlüssel (recv.key). Zum Senden "
                 "brauchst du deinen send.key.")
    offset = key["send_offset"]
    # Pro Nachricht: len(plaintext) Bytes Keystream + 32 Bytes MAC-Schlüssel.
    need = len(plaintext) + MAC_KEY_LEN
    segment = key["material"][offset:offset + need]
    if len(segment) < need:
        sys.exit(f"[!] Sendeschlüssel erschöpft: ab Offset {offset} bleiben nur "
                 f"{len(segment)} Bytes, Nachricht+MAC brauchen {need}. "
                 "Erzeugt ein neues Schlüsselpaar.")

    ciphertext = otp_xor(plaintext, key["material"], offset)
    mac_key = key["material"][offset + len(plaintext):offset + need]
    pre_header = struct.pack(STEGO_PRE_FMT, STEGO_MAGIC, STEGO_VERSION,
                             key["stream_id"], offset, len(ciphertext))
    tag = poly1305_mac(pre_header + ciphertext, mac_key)
    embed(args.carrier, pre_header + tag + ciphertext, out_path)
    advance_send_offset(args.key, offset + need)

    print(f"[+] Verschlüsselt + authentifiziert ({len(plaintext)} B) ab "
          f"Sende-Offset {offset}.")
    print(f"[+] Versteckt in: {out_path}")
    print(f"[+] Sende-Offset im Schlüssel weitergezählt -> {offset + need} "
          f"(Nachricht {len(plaintext)} + MAC {MAC_KEY_LEN}).")
    print("[i] Als DATEI/Dokument verschicken, nicht als Foto!")


def cmd_reveal(args):
    info = extract(args.stego)
    key = read_key(args.key)
    if key["role"] != ROLE_RECV:
        sys.exit("[!] Das ist ein Sendeschlüssel (send.key). Zum Empfangen "
                 "brauchst du deinen recv.key.")
    if info["stream_id"] != key["stream_id"]:
        sys.exit("[!] Falscher Schlüssel: dieses Bild gehört zu einem anderen "
                 "Kontakt/Strom. Nimm den passenden recv.key.")
    offset, ct = info["offset"], info["ciphertext"]

    if info["version"] == STEGO_VERSION:
        # Authentifiziert: MAC-Schlüssel liegt direkt hinter dem Keystream.
        need = len(ct) + MAC_KEY_LEN
        segment = key["material"][offset:offset + need]
        if len(segment) < need:
            sys.exit(f"[!] Schlüssel passt nicht: ab Offset {offset} brauche ich "
                     f"{need} Bytes (Nachricht+MAC), habe nur {len(segment)}.")
        mac_key = key["material"][offset + len(ct):offset + need]
        expected = poly1305_mac(info["mac_input_header"] + ct, mac_key)
        if not ct_equal(expected, info["tag"]):
            sys.exit("[!] Authentifizierung fehlgeschlagen: das Bild wurde "
                     "verändert, oder der Schlüssel ist nicht synchron "
                     "(Offset-Drift). Nachricht NICHT vertrauen.")
    else:
        segment = key["material"][offset:offset + len(ct)]
        if len(segment) < len(ct):
            sys.exit(f"[!] Schlüssel passt nicht: ab Offset {offset} brauche ich "
                     f"{len(ct)} Bytes, habe nur {len(segment)}.")
        print("[i] Warnung: altes SPS2-Bild — unauthentifiziert, "
              "keine Manipulationssicherheit.", file=sys.stderr)

    plaintext = otp_xor(ct, key["material"], offset)

    if args.outfile:
        with open(args.outfile, "wb") as fh:
            fh.write(plaintext)
        print(f"[+] Entschlüsselt ({len(plaintext)} B, Offset {offset}) "
              f"-> {args.outfile}")
    else:
        try:
            sys.stdout.write(plaintext.decode("utf-8"))
            if not plaintext.endswith(b"\n"):
                sys.stdout.write("\n")
        except UnicodeDecodeError:
            sys.stdout.buffer.write(plaintext)


def build_parser():
    p = argparse.ArgumentParser(
        prog="schattenpost",
        description="One-Time-Pad + LSB-Steganografie, richtungsgetrennte Schlüssel.",
    )
    sub = p.add_subparsers(dest="cmd", required=True)

    g = sub.add_parser("genkey", help="Gerichtetes Schlüsselpaar erzeugen.")
    g.add_argument("-o", "--output", required=True,
                   help="Basisname (erzeugt <name>_meine_seite/ + _partner_seite/).")
    g.add_argument("-s", "--size", type=int, default=1_000_000,
                   help="Bytes je Richtung (Default 1 MB).")
    g.add_argument("--force", action="store_true")
    g.set_defaults(func=cmd_genkey)

    h = sub.add_parser("hide", help="Nachricht verschlüsseln + verstecken.")
    h.add_argument("-t", "--text")
    h.add_argument("-i", "--infile")
    h.add_argument("-k", "--key", required=True, help="Dein send.key.")
    h.add_argument("-c", "--carrier", required=True, help="Trägerbild.")
    h.add_argument("-o", "--output",
                   help="Stego-PNG (Default: unauffälliger Screenshot_-Name "
                        "mit Zeitstempel).")
    h.set_defaults(func=cmd_hide)

    r = sub.add_parser("reveal", help="Nachricht extrahieren + entschlüsseln.")
    r.add_argument("-s", "--stego", required=True)
    r.add_argument("-k", "--key", required=True, help="Dein recv.key.")
    r.add_argument("-o", "--outfile")
    r.set_defaults(func=cmd_reveal)
    return p


def main():
    args = build_parser().parse_args()
    try:
        args.func(args)
    except (ValueError, FileNotFoundError, struct.error) as exc:
        sys.exit(f"[!] {exc}")


if __name__ == "__main__":
    main()