scripts/generate_key.py

#!/usr/bin/env python3
"""Provision device credentials into Redis for imageO_v3."""

from __future__ import annotations

import argparse
import os
import secrets
import string
import sys
from dataclasses import asdict, dataclass
from typing import List

import redis
from dotenv import load_dotenv


@dataclass
class DeviceCredentials:
    """Single device credential record."""

    device_id: str
    api_key: str
    secret: str
    is_active: bool


class DeviceProvisioner:
    """Generate unique device credentials and persist to Redis."""

    DEVICE_ID_PREFIX = "DEV"
    API_KEY_PREFIX = "sk"
    API_KEY_LENGTH = 32
    SECRET_LENGTH = 64

    def __init__(self, redis_client: redis.Redis):
        self.redis_client = redis_client
        self.generated_ids_this_run = set()
        self.generated_api_keys_this_run = set()
        self.generated_secrets_this_run = set()

    @staticmethod
    def generate_random_string(length: int, charset: str | None = None) -> str:
        if charset is None:
            charset = string.ascii_uppercase + string.digits
        return "".join(secrets.choice(charset) for _ in range(length))

    def generate_device_id(self) -> str:
        while True:
            random_part = self.generate_random_string(12)
            candidate = (
                f"{self.DEVICE_ID_PREFIX}-{random_part[0:4]}-"
                f"{random_part[4:8]}-{random_part[8:12]}"
            )
            redis_key = f"device:{candidate}"
            if candidate not in self.generated_ids_this_run and not self.redis_client.exists(redis_key):
                self.generated_ids_this_run.add(candidate)
                return candidate

    def generate_api_key(self) -> str:
        while True:
            api_key = f"{self.API_KEY_PREFIX}_{secrets.token_hex(self.API_KEY_LENGTH // 2)}"
            if api_key not in self.generated_api_keys_this_run:
                self.generated_api_keys_this_run.add(api_key)
                return api_key

    def generate_secret(self) -> str:
        while True:
            secret = secrets.token_hex(self.SECRET_LENGTH // 2)
            if secret not in self.generated_secrets_this_run:
                self.generated_secrets_this_run.add(secret)
                return secret

    def provision_device(self, is_active: bool = True) -> DeviceCredentials:
        return DeviceCredentials(
            device_id=self.generate_device_id(),
            api_key=self.generate_api_key(),
            secret=self.generate_secret(),
            is_active=is_active,
        )

    def provision_multiple(self, count: int) -> List[DeviceCredentials]:
        return [self.provision_device() for _ in range(count)]


def main() -> None:
    """CLI entrypoint."""
    load_dotenv()
    parser = argparse.ArgumentParser(
        description="Generate and store secure device credentials in Redis.",
        epilog="Example: python scripts/generate_key.py --count 10",
    )
    parser.add_argument("-c", "--count", type=int, required=True, help="Number of devices to provision")
    args = parser.parse_args()

    redis_url = os.getenv("REDIS_URL")
    if not redis_url:
        print("Error: REDIS_URL not found in environment or .env file.", file=sys.stderr)
        raise SystemExit(1)

    try:
        redis_client = redis.from_url(redis_url, decode_responses=True)
        redis_client.ping()
    except redis.exceptions.RedisError as exc:
        print(f"Error: Could not connect to Redis. Reason: {exc}", file=sys.stderr)
        raise SystemExit(1) from exc

    provisioner = DeviceProvisioner(redis_client)
    devices = provisioner.provision_multiple(args.count)

    try:
        with redis_client.pipeline() as pipe:
            for device in devices:
                redis_key = f"device:{device.device_id}"
                payload = asdict(device)
                payload.pop("device_id", None)
                payload["is_active"] = "1" if payload["is_active"] else "0"
                pipe.hset(redis_key, mapping=payload)
            pipe.execute()
    except redis.exceptions.RedisError as exc:
        print(f"Error: Failed to write to Redis. Reason: {exc}", file=sys.stderr)
        raise SystemExit(1) from exc

    print(f"Provisioned {len(devices)} device(s).")
    print("--- Newly Provisioned Devices ---")
    for device in devices:
        print(f"Device ID: {device.device_id}")
    print("---------------------------------")


if __name__ == "__main__":
    main()