ria-ran--sim-drop/server/core_adapters/open5gs.py

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"):
        # 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]

    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": "internet",
                            "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}))