#!/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()