#!/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.

WICHTIG:
  * Verschickt das Stego-Bild IMMER als Datei/Dokument, niemals als "Foto" —
    Messenger jagen Fotos durch verlustbehaftete Kompression und zerstören dabei
    die versteckten Bits. Als Dokument bleibt das PNG unangetastet.
  * Das Pad (der Schlüssel) muss auf beiden Seiten byte-identisch sein und darf
    NIE zweimal für dieselben Bytes benutzt werden. Der Offset-Zähler
    (pad.key.offset) sorgt automatisch dafür, dass kein Byte doppelt verbraucht
    wird — solange ihr ihn nicht manuell zurückdreht.

Subcommands:
  genkey   Erzeugt ein echtes Zufalls-Pad.
  hide     Verschlüsselt + versteckt eine Nachricht in einem Trägerbild.
  reveal   Extrahiert + entschlüsselt eine Nachricht aus einem Stego-Bild.

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

import argparse
import os
import struct
import sys

from PIL import Image

MAGIC = b"SPST"          # Schattenpost-Container-Kennung
VERSION = 1
# Header: magic(4) + version(1) + offset(8, uint64 BE) + length(4, uint32 BE)
HEADER_FMT = ">4sBQI"
HEADER_LEN = struct.calcsize(HEADER_FMT)


# --------------------------------------------------------------------------- #
# One-Time-Pad
# --------------------------------------------------------------------------- #
def otp_xor(data: bytes, key: bytes) -> bytes:
    """XOR data gegen key. key muss mindestens so lang wie data sein."""
    if len(key) < len(data):
        raise ValueError(
            f"Pad zu kurz: brauche {len(data)} Bytes, habe nur {len(key)}."
        )
    return bytes(d ^ k for d, k in zip(data, key))


def load_pad(path: str) -> bytes:
    with open(path, "rb") as fh:
        return fh.read()


def read_offset(pad_path: str) -> int:
    """Nächster unbenutzter Offset im Pad. Fehlt die Datei -> 0."""
    off_path = pad_path + ".offset"
    if not os.path.exists(off_path):
        return 0
    with open(off_path, "r") as fh:
        return int(fh.read().strip() or "0")


def write_offset(pad_path: str, value: int) -> None:
    with open(pad_path + ".offset", "w") as fh:
        fh.write(str(value))


# --------------------------------------------------------------------------- #
# LSB-Steganografie
# --------------------------------------------------------------------------- #
def _bits_from_bytes(data: bytes):
    """Liefert die Bits von data, MSB zuerst pro Byte."""
    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:
    """Bettet payload in die LSBs von carrier ein und speichert als PNG."""
    img = Image.open(carrier_path).convert("RGB")
    width, height = img.size
    capacity_bits = width * height * 3
    need_bits = len(payload) * 8
    if need_bits > capacity_bits:
        raise ValueError(
            f"Trägerbild zu klein: brauche {need_bits} Bits, "
            f"Platz für {capacity_bits} ({need_bits // 8} vs "
            f"{capacity_bits // 8} Bytes). 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]
            channels = [r, g, b]
            for c in range(3):
                bit = next(bits, None)
                if bit is None:
                    done = True
                    break
                channels[c] = (channels[c] & ~1) | bit
            px[x, y] = (channels[0], channels[1], channels[2])
            if done:
                break

    img.save(out_path, "PNG")


def _read_bits(px, width, height, n_bits):
    """Liest n_bits LSBs linear aus dem Bild und gibt sie als Bytes zurück."""
    out = bytearray()
    acc = 0
    count = 0
    read = 0
    for y in range(height):
        for x in range(width):
            r, g, b = px[x, y]
            for c in (r, g, b):
                acc = (acc << 1) | (c & 1)
                count += 1
                read += 1
                if count == 8:
                    out.append(acc)
                    acc = 0
                    count = 0
                if read >= n_bits:
                    return bytes(out)
    return bytes(out)


def extract(stego_path: str) -> tuple[int, bytes]:
    """Liest Header + Ciphertext aus einem Stego-Bild. Gibt (offset, ct)."""
    img = Image.open(stego_path).convert("RGB")
    width, height = img.size
    px = img.load()

    header = _read_bits(px, width, height, HEADER_LEN * 8)
    magic, version, offset, length = struct.unpack(HEADER_FMT, header)
    if magic != MAGIC:
        raise ValueError(
            "Kein Schattenpost-Container gefunden (Magic passt nicht). "
            "Wurde das Bild vielleicht doch als Foto komprimiert?"
        )
    if version != VERSION:
        raise ValueError(f"Unbekannte Version {version}.")

    total_bits = (HEADER_LEN + length) * 8
    blob = _read_bits(px, width, height, total_bits)
    ciphertext = blob[HEADER_LEN:HEADER_LEN + length]
    return offset, ciphertext


# --------------------------------------------------------------------------- #
# Commands
# --------------------------------------------------------------------------- #
def cmd_genkey(args):
    size = args.size
    if os.path.exists(args.output) and not args.force:
        sys.exit(f"[!] {args.output} existiert schon. --force zum Überschreiben.")
    with open(args.output, "wb") as fh:
        fh.write(os.urandom(size))
    write_offset(args.output, 0)
    print(f"[+] Pad erzeugt: {args.output} ({size} Bytes, echtes os.urandom)")
    print(f"[+] Offset-Zähler zurückgesetzt: {args.output}.offset -> 0")
    print("[i] Gib dieses Pad sicher an deinen Gesprächspartner weiter "
          "(z.B. USB-Stick, persönlich). Beide brauchen dieselbe Datei.")


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.")

    pad = load_pad(args.key)
    offset = read_offset(args.key)
    segment = pad[offset:offset + len(plaintext)]
    if len(segment) < len(plaintext):
        sys.exit(
            f"[!] Pad erschöpft: ab Offset {offset} bleiben nur "
            f"{len(segment)} Bytes, Nachricht braucht {len(plaintext)}. "
            "Erzeuge/tausche ein neues Pad."
        )

    ciphertext = otp_xor(plaintext, segment)
    header = struct.pack(HEADER_FMT, MAGIC, VERSION, offset, len(ciphertext))
    payload = header + ciphertext

    embed(args.carrier, payload, args.output)
    write_offset(args.key, offset + len(plaintext))

    print(f"[+] Nachricht verschlüsselt ({len(plaintext)} Bytes) "
          f"mit Pad-Offset {offset}.")
    print(f"[+] Versteckt in: {args.output}")
    print(f"[+] Pad-Offset weitergezählt -> {offset + len(plaintext)}")
    print("[i] Verschicke das PNG als DATEI/Dokument, nicht als Foto!")


def cmd_reveal(args):
    offset, ciphertext = extract(args.stego)
    pad = load_pad(args.key)
    segment = pad[offset:offset + len(ciphertext)]
    if len(segment) < len(ciphertext):
        sys.exit(
            f"[!] Pad passt nicht: ab Offset {offset} brauche ich "
            f"{len(ciphertext)} Bytes, habe nur {len(segment)}."
        )
    plaintext = otp_xor(ciphertext, segment)

    if args.outfile:
        with open(args.outfile, "wb") as fh:
            fh.write(plaintext)
        print(f"[+] Entschlüsselt ({len(plaintext)} Bytes, 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 für vertrauliche Nachrichten.",
    )
    sub = p.add_subparsers(dest="cmd", required=True)

    g = sub.add_parser("genkey", help="Echtes Zufalls-Pad erzeugen.")
    g.add_argument("-o", "--output", required=True, help="Pad-Datei.")
    g.add_argument("-s", "--size", type=int, default=1_000_000,
                   help="Größe in Bytes (Default 1 MB).")
    g.add_argument("--force", action="store_true", help="Überschreiben erlauben.")
    g.set_defaults(func=cmd_genkey)

    h = sub.add_parser("hide", help="Nachricht verschlüsseln + verstecken.")
    h.add_argument("-t", "--text", help="Nachrichtentext (sonst --infile/stdin).")
    h.add_argument("-i", "--infile", help="Datei als Nachricht.")
    h.add_argument("-k", "--key", required=True, help="Pad-Datei.")
    h.add_argument("-c", "--carrier", required=True, help="Trägerbild.")
    h.add_argument("-o", "--output", required=True, help="Stego-PNG-Ausgabe.")
    h.set_defaults(func=cmd_hide)

    r = sub.add_parser("reveal", help="Nachricht extrahieren + entschlüsseln.")
    r.add_argument("-s", "--stego", required=True, help="Stego-PNG.")
    r.add_argument("-k", "--key", required=True, help="Pad-Datei.")
    r.add_argument("-o", "--outfile", help="Ausgabedatei (sonst stdout).")
    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()