#!/usr/bin/env bash
set -euo pipefail

IMAGE_NAME="kyber-hybrid:latest"
CONTAINER_NAME="kyber-hybrid-demo"

# Create a build context
WORKDIR="$(pwd)/kyber_hybrid_build"
mkdir -p "$WORKDIR"
cd "$WORKDIR"

# -------------------------------------------------------------------
# Dockerfile
# -------------------------------------------------------------------
cat > Dockerfile <<'EOF'
FROM ubuntu:24.04

ENV DEBIAN_FRONTEND=noninteractive
ENV TZ=UTC

# 1) base system and build deps
RUN apt-get update && apt-get install -y --no-install-recommends \
    ca-certificates curl git build-essential cmake ninja-build pkg-config \
    python3 python3-pip python3-dev python3-venv \
    libssl-dev openssl \
    && rm -rf /var/lib/apt/lists/*

# 2) build and install liboqs
#    note: we use a specific commit to avoid breaking changes over time
#    you can update this pin if you prefer a newer release from OQS
ARG OQS_COMMIT=main
RUN git clone --depth 1 --branch ${OQS_COMMIT} https://github.com/open-quantum-safe/liboqs.git /tmp/liboqs && \
    mkdir -p /tmp/liboqs/build && cd /tmp/liboqs/build && \
    cmake -GNinja .. -DOQS_BUILD_ONLY_LIB=ON -DCMAKE_INSTALL_PREFIX=/usr/local \
      -DOQS_USE_OPENSSL=ON -DBUILD_SHARED_LIBS=ON && \
    ninja && ninja install && \
    ldconfig && \
    rm -rf /tmp/liboqs

# 3) python packages
#    oqs will link against the shared lib we just installed
RUN python3 -m pip install --no-cache-dir --upgrade pip && \
    python3 -m pip install --no-cache-dir oqs cryptography

# keep runtime small by removing build tools if you like
# (left installed here to ease debugging; uncomment to slim)
# RUN apt-get purge -y git build-essential cmake ninja-build pkg-config python3-dev libssl-dev && \
#     apt-get autoremove -y && \
#     apt-get clean -y

# 4) app
WORKDIR /app
COPY kyber_hybrid.py /app/kyber_hybrid.py

# default command runs the demo script
CMD ["python3", "/app/kyber_hybrid.py"]
EOF

# -------------------------------------------------------------------
# kyber_hybrid.py
# -------------------------------------------------------------------
cat > kyber_hybrid.py <<'EOF'
import base64
import json
import os

# pip: oqs, cryptography
import oqs
from cryptography.hazmat.primitives.ciphers.aead import AESGCM

def b64(x: bytes) -> str:
    return base64.b64encode(x).decode()

def ub64(s: str) -> bytes:
    return base64.b64decode(s.encode())

def generate_kyber_keypair(kyber_alg: str = "Kyber512"):
    with oqs.KeyEncapsulation(kyber_alg) as kem:
        pk = kem.generate_keypair()
        sk = kem.export_secret_key()
    return pk, sk

def encapsulate_key(public_key: bytes, kyber_alg: str = "Kyber512"):
    with oqs.KeyEncapsulation(kyber_alg) as kem:
        ct, ss = kem.encap_secret(public_key)
    return ct, ss

def decapsulate_key(encapsulated_ciphertext: bytes, private_key: bytes, kyber_alg: str = "Kyber512"):
    with oqs.KeyEncapsulation(kyber_alg) as kem:
        kem.load_secret_key(private_key)
        shared_secret = kem.decap_secret(encapsulated_ciphertext)
    return shared_secret

def aesgcm_encrypt(plaintext: bytes, key: bytes):
    aesgcm = AESGCM(key)
    nonce = os.urandom(12)
    ct = aesgcm.encrypt(nonce, plaintext, associated_data=None)
    return nonce, ct

def aesgcm_decrypt(nonce: bytes, ct: bytes, key: bytes):
    aesgcm = AESGCM(key)
    pt = aesgcm.decrypt(nonce, ct, associated_data=None)
    return pt

def hybrid_encrypt(plaintext: bytes, public_key: bytes, kyber_alg: str = "Kyber512"):
    kem_ct, shared_secret = encapsulate_key(public_key, kyber_alg)
    # shared_secret length is suitable for AESGCM as a 256-bit key
    nonce, ct = aesgcm_encrypt(plaintext, shared_secret)
    payload = {
        "kem_ct": b64(kem_ct),
        "aes_nonce": b64(nonce),
        "aes_ct": b64(ct)
    }
    return payload

def hybrid_decrypt(payload: dict, private_key: bytes, kyber_alg: str = "Kyber512"):
    kem_ct = ub64(payload["kem_ct"])
    nonce = ub64(payload["aes_nonce"])
    aes_ct = ub64(payload["aes_ct"])
    shared_secret = decapsulate_key(kem_ct, private_key, kyber_alg)
    pt = aesgcm_decrypt(nonce, aes_ct, shared_secret)
    return pt

if __name__ == "__main__":
    # 1) keypair
    public_key, private_key = generate_kyber_keypair("Kyber512")
    print("Public key size:", len(public_key))
    print("Private key size:", len(private_key))

    # 2) example secret document
    secret_message = b"This is a very secret document to store in Qdrant."

    # 3) encrypt
    enc_payload = hybrid_encrypt(secret_message, public_key, "Kyber512")
    print("Encrypted payload (json):")
    print(json.dumps(enc_payload, indent=2))

    # 4) decrypt
    recovered = hybrid_decrypt(enc_payload, private_key, "Kyber512")
    print("Recovered:", recovered.decode())
EOF

# -------------------------------------------------------------------
# Build image
# -------------------------------------------------------------------
echo "Building Docker image: ${IMAGE_NAME}"
docker build -t "${IMAGE_NAME}" .

# Clean up any prior container with the same name
if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}\$"; then
  docker rm -f "${CONTAINER_NAME}" >/dev/null 2>&1 || true
fi

# -------------------------------------------------------------------
# Run container once to demonstrate
# -------------------------------------------------------------------
echo "Running demo container..."
docker run --name "${CONTAINER_NAME}" --rm "${IMAGE_NAME}"