#!/usr/bin/env python3
"""Camera-based entropy source.

Fetches MJPEG frames from a network camera, hashes them, and feeds the
hash into a NIST SP 800-90A HMAC-DRBG. Output is the DRBG's generated
bytes, so the camera genuinely contributes entropy (unlike the previous
version which computed a hash and then discarded it).
"""
import argparse
import hashlib
import hmac
import sys
import time

import cv2

from health import check as health_check, HealthCheckError

sys.stdout.reconfigure(line_buffering=True)


class HMACDRBG:
    """NIST SP 800-90A HMAC-DRBG using SHA-256.

    Seeded with camera frame entropy. Supports reseed so each frame
    mixes fresh entropy into the internal state.
    """

    def __init__(self, entropy, personalization=b""):
        self.K = b"\x00" * 32
        self.V = b"\x01" * 32
        self.reseed_counter = 0
        self._update(entropy + personalization)

    def _update(self, provided_data):
        self.K = hmac.new(self.K, self.V + b"\x00" + provided_data, hashlib.sha256).digest()
        self.V = hmac.new(self.K, self.V, hashlib.sha256).digest()
        self.K = hmac.new(self.K, self.V + b"\x01" + provided_data, hashlib.sha256).digest()
        self.V = hmac.new(self.K, self.V, hashlib.sha256).digest()

    def reseed(self, entropy):
        self._update(entropy)
        self.reseed_counter = 0

    def generate(self, num_bytes):
        if self.reseed_counter > 10000:
            raise RuntimeError("DRBG requires reseed before generating more output.")
        self.reseed_counter += 1
        out = b""
        while len(out) < num_bytes:
            self.V = hmac.new(self.K, self.V, hashlib.sha256).digest()
            out += self.V
        self._update(b"")
        return out[:num_bytes]


def fetch_frame_opencv(url):
    cap = cv2.VideoCapture(url)
    ret, frame = cap.read()
    cap.release()
    if not ret:
        raise ValueError("Failed to fetch frame.")
    return frame


def frame_entropy(frame):
    """Hash a captured frame to a 32-byte digest."""
    _, img_encoded = cv2.imencode(".jpg", frame)
    return hashlib.sha256(img_encoded.tobytes()).digest()


def frame_raw_bytes(frame):
    """Return the raw JPEG-encoded frame bytes (for health checks)."""
    _, img_encoded = cv2.imencode(".jpg", frame)
    return img_encoded.tobytes()


def emit(blob, out_path):
    if out_path:
        with open(out_path, "ab") as f:
            f.write(blob)
    else:
        sys.stdout.write(blob.hex() + "\n")


def main():
    parser = argparse.ArgumentParser(description="Generate random numbers from camera entropy.")
    parser.add_argument("--url", default="http://192.168.0.200/mjpg/video.mjpg", help="Camera MJPEG stream URL")
    parser.add_argument("--single", type=int, default=1, help="Generate N random numbers and exit (default: 1)")
    parser.add_argument("--bit", type=int, default=256, help="Bit length of each random number (default: 256)")
    parser.add_argument("--loop", action="store_true", help="Run forever, generating numbers every second")
    parser.add_argument("--interval", type=float, default=1.0, help="Seconds between iterations in --loop (default: 1.0)")
    parser.add_argument("--out", help="Append raw bytes to this path (e.g. a FIFO) instead of printing hex")
    parser.add_argument("--no-health", action="store_true", help="Disable FIPS 140-2 health checks on raw frames")
    args = parser.parse_args()

    if args.bit <= 0 or args.bit % 8 != 0:
        print("Error: --bit must be a positive multiple of 8", file=sys.stderr)
        return 2
    num_bytes = args.bit // 8

    if args.loop and args.single != 1:
        print("Error: --loop and --single are mutually exclusive", file=sys.stderr)
        return 2

    def produce():
        frame = fetch_frame_opencv(args.url)
        raw = frame_raw_bytes(frame)
        if not args.no_health:
            health_check(raw)
        entropy = hashlib.sha256(raw).digest()
        drbg = HMACDRBG(entropy)
        return drbg.generate(num_bytes)

    if args.loop:
        while True:
            try:
                blob = produce()
                emit(blob, args.out)
            except Exception as e:
                print(f"Error: {e}", file=sys.stderr)
            time.sleep(args.interval)
    else:
        for i in range(args.single):
            try:
                blob = produce()
                emit(blob, args.out)
            except Exception as e:
                print(f"Error: {e}", file=sys.stderr)
                return 1
            if args.single > 1 and i < args.single - 1:
                time.sleep(args.interval)
    return 0


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