import logging from pymongo import MongoClient from .base import CoreAdapter logger = logging.getLogger(__name__) class Open5GSAdapter(CoreAdapter): def __init__(self, mongodb_uri: str, database_name: str = "open5gs", apn: str = "internet"): # Short timeouts so a down/unreachable core fails the provision fast # instead of blocking the request for pymongo's 30s default. self.client = MongoClient( mongodb_uri, serverSelectionTimeoutMS=3000, connectTimeoutMS=3000, socketTimeoutMS=3000, ) self.db = self.client[database_name] self.apn = apn def add_subscriber(self, imsi: str, ki: str, opc: str) -> None: """Upsert a subscriber into the Open5GS MongoDB collection.""" doc = { "imsi": imsi, "security": { "k": ki, "opc": opc, "amf": "8000", "sqn": 0, }, "ambr": { "downlink": {"value": 1, "unit": 3}, "uplink": {"value": 1, "unit": 3}, }, "slice": [ { "sst": 1, "default_indicator": True, "session": [ { "name": self.apn, "type": 3, "ambr": { "downlink": {"value": 1, "unit": 3}, "uplink": {"value": 1, "unit": 3}, }, } ], } ], } self.db.subscribers.replace_one({"imsi": imsi}, doc, upsert=True) logger.info(f"Added subscriber {imsi} to Open5GS") def remove_subscriber(self, imsi: str) -> None: self.db.subscribers.delete_one({"imsi": imsi}) logger.info(f"Removed subscriber {imsi} from Open5GS") def list_subscribers(self) -> list: return list(self.db.subscribers.find({}, {"_id": 0}))