Spaces:
Sleeping
Sleeping
File size: 4,460 Bytes
ee54f79 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | #!/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()
|