From 72374e439a809c3b04ee073a78e6caa196b53ed9 Mon Sep 17 00:00:00 2001 From: Matthew Johnson Date: Sat, 24 Apr 2021 23:25:51 -0700 Subject: [PATCH] Some cleanup for the db migration script (#956) --- misc/db/python/Open5GS.py | 7 +- misc/db/python/README.md | 2 +- misc/db/python/SchemaUpdater.py | 147 +++++++++++++++++---------- misc/db/python/test_SchemaUpdater.py | 93 +++++++++++++++++ 4 files changed, 186 insertions(+), 63 deletions(-) create mode 100644 misc/db/python/test_SchemaUpdater.py diff --git a/misc/db/python/Open5GS.py b/misc/db/python/Open5GS.py index 098ca2e45..23b4c103e 100644 --- a/misc/db/python/Open5GS.py +++ b/misc/db/python/Open5GS.py @@ -1,4 +1,3 @@ -import mongo import pymongo import random import bson @@ -8,7 +7,6 @@ class Open5GS: self.server = server self.port = port - def GetSubscribers(self): myclient = pymongo.MongoClient("mongodb://" + str(self.server) + ":" + str(self.port) + "/") mydb = myclient["open5gs"] @@ -30,14 +28,12 @@ class Open5GS: for x in mydoc: print(x) return x - def AddSubscriber(self, sub_data): - myclient = pymongo.MongoClient("mongodb://" + str(self.server) + ":" + str(self.port) + "/") mydb = myclient["open5gs"] mycol = mydb["subscribers"] - + x = mycol.insert_one(sub_data) print("Added subscriber with Inserted ID : " + str(x.inserted_id)) return x.inserted_id @@ -53,7 +49,6 @@ class Open5GS: print(x) return True - def DeleteSubscriber(self, imsi): myclient = pymongo.MongoClient("mongodb://" + str(self.server) + ":" + str(self.port) + "/") mydb = myclient["open5gs"] diff --git a/misc/db/python/README.md b/misc/db/python/README.md index cfcee19ce..9ba1b76e0 100644 --- a/misc/db/python/README.md +++ b/misc/db/python/README.md @@ -1,6 +1,6 @@ ## Open5GS Python Library -Basic Python library to interface with MongoDB subscriber DB in Open5GS HSS / PCRF. Requires Python 3+, mongo, pymongo and bson. (All available through PIP) +Basic Python library to interface with MongoDB subscriber DB in Open5GS HSS / PCRF. Requires Python 3+, pymongo and bson. (All available through PIP) If you are planning to run this on a different machine other than localhost (the machine hosting the MongoDB service) you will need to enable remote access to MongoDB by binding it's IP to 0.0.0.0: diff --git a/misc/db/python/SchemaUpdater.py b/misc/db/python/SchemaUpdater.py index 3b49fd2cb..03a21a774 100644 --- a/misc/db/python/SchemaUpdater.py +++ b/misc/db/python/SchemaUpdater.py @@ -3,66 +3,101 @@ #Additional functionlality like PCC rules, static assignment etc, not tested. If it's not listed below it's probably not migrated by this script. #Written by @nickvsnetworking 30/03/2021 -import json -import sys -import random, string -import mongo +import copy import pymongo -myclient = pymongo.MongoClient("mongodb://localhost:27017/") -mydb = myclient["open5gs"] -mycol = mydb["subscribers"] -subs_list = [] -for x in mycol.find(): - if 'schema_version' not in x: - print("Subscriber record " + str(x['imsi']) + " needs updating") - old_template_json = x - print(old_template_json) - #Set AMBR Values to new format (Old format is in bits per second) - try: - uplink = old_template_json['ambr']['uplink'] - old_template_json['ambr']['uplink'] = {} - old_template_json['ambr']['uplink']['value'] = uplink - old_template_json['ambr']['uplink']['unit'] = 0 - except Exception as e: - print(e) - print("Failed to set Uplink AMBR values") - try: - downlink = old_template_json['ambr']['downlink'] - old_template_json['ambr']['downlink'] = {} - old_template_json['ambr']['downlink']['value'] = downlink - old_template_json['ambr']['downlink']['unit'] = 0 - except Exception as e: - print(e) - print("Failed to set Downlink AMBR values") +def migrate_all_subscribers(mycol): + """Migrates all subscribers in the mycol collection from schema version 0 to version 1 + """ + for x in mycol.find(): + if 'schema_version' not in x: + imsi = x['imsi'] + print("Subscriber record " + str(imsi) + " needs updating") - #Propogate APN / DDN Slice Details - old_template_json['slice'] = [] - old_template_json['slice'].append({"sst": 1, "default_indicator" : True, "session" : []}) - - i = 0 - while i < len(old_template_json['pdn']): - ddn_dict = {} - ddn_dict['name'] = old_template_json['pdn'][i]['apn'] - ddn_dict['type'] = old_template_json['pdn'][i]['type'] - ddn_dict['pcc_rule'] = old_template_json['pdn'][i]['pcc_rule'] - ddn_dict['qos'] = old_template_json['pdn'][i]['qos'] - ddn_dict['qos']['index'] = old_template_json['pdn'][i]['qos']['qci'] - ddn_dict['qos']['arp'] = old_template_json['pdn'][i]['qos']['arp'] - ddn_dict['ambr'] = {"uplink": {"value": old_template_json['pdn'][i]['ambr']['uplink'], "unit": 0}, "downlink": {"value": old_template_json['pdn'][i]['ambr']['downlink'], "unit": 0}} - i += 1 - old_template_json['slice'][0]['session'].append(ddn_dict) - - #Remove old PDN info - #del old_template_json['pdn'] + print("Current value:", x) + new_subscriber = create_v1_from_v0(x) + print("Migrated value:", new_subscriber) - #Add "schema_version" feild - old_template_json['schema_version'] = 1 + #Write back to MongoDB + myquery = { "imsi": str(imsi) } + newvalues = { + "$set": new_subscriber, + "$unset": {"pdn": 1} + } + mycol.update_one(myquery, newvalues) + print("Updated OK") - #Write back to MongoDB - myquery = { "imsi": str(old_template_json['imsi'])} - newvalues = { "$set": old_template_json } - mycol.update_one(myquery, newvalues) - print("Updated OK") +def create_v1_from_v0(old_sub): + """Create a v1 subscriber from an existing v0 subscriber + """ + # Make a copy to avoid mutating the existing subscriber object so it can be + # re-used for other parts of the migration. + new_sub = copy.deepcopy(old_sub) + + # Remove old PDN info + del new_sub['pdn'] + + # Set AMBR Values to new format (Old format is in bits per second) + new_sub['ambr']['uplink'] = {} + new_sub['ambr']['uplink']['value'] = old_sub['ambr']['uplink'] + new_sub['ambr']['uplink']['unit'] = 0 + + new_sub['ambr']['downlink'] = {} + new_sub['ambr']['downlink']['value'] = old_sub['ambr']['downlink'] + new_sub['ambr']['downlink']['unit'] = 0 + + #Propogate APN / DDN Slice Details + new_sub['slice'] = [] + new_sub['slice'].append({"sst": 1, "default_indicator" : True, "session" : []}) + + for pdn_entry in old_sub["pdn"]: + session = _create_session_from_pdn(pdn_entry) + new_sub['slice'][0]['session'].append(session) + + #Add "schema_version" feild + new_sub['schema_version'] = 1 + + return new_sub + + +def _create_session_from_pdn(pdn): + """Builds a new session object from an existing PDN""" + session = {} + session['name'] = pdn['apn'] + session['type'] = pdn['type'] + session['ambr'] = { + "uplink": { + "value": pdn['ambr']['uplink'], + "unit": 0 + }, + "downlink": { + "value": pdn['ambr']['downlink'], + "unit": 0 + } + } + + if "qos" in pdn: + session["qos"] = { + "index": pdn["qos"]["qci"], + "arp": pdn["qos"]["arp"] + } + if "smf" in pdn: + session["smf"] = pdn["smf"] + if "ue" in pdn: + session["ue"] = pdn["ue"] + + if ("pcc_rule" in pdn) and (len(pdn['pcc_rule']) != 0): + raise NotImplementedError("PCC Rule Migration Not Implemented") + else: + session["pcc_rule"] = [] + + return session + + +if __name__ == "__main__": + myclient = pymongo.MongoClient("mongodb://localhost:27017/") + mydb = myclient["open5gs"] + + migrate_all_subscribers(mycol=mydb["subscribers"]) diff --git a/misc/db/python/test_SchemaUpdater.py b/misc/db/python/test_SchemaUpdater.py new file mode 100644 index 000000000..6506e64f1 --- /dev/null +++ b/misc/db/python/test_SchemaUpdater.py @@ -0,0 +1,93 @@ +import unittest +import pymongo + +import SchemaUpdater + +class TestSchemaUpdater(unittest.TestCase): + def setUp(self): + self.legacy_sub = { + 'imsi': '999990000000001', + 'access_restriction_data': 32, + 'ambr': { + 'downlink': 1024000, + 'uplink': 1024000 + }, + 'network_access_mode': 2, + 'pdn': [ + { + 'apn': 'internet', + 'pcc_rule': [], + 'ambr': { + 'downlink': 1024000, + 'uplink': 1024000 + }, + 'qos': { + 'qci': 9, + 'arp': { + 'priority_level': 8, + 'pre_emption_vulnerability': 1, + 'pre_emption_capability': 0 + } + }, + 'type': 0, + 'ue': { + 'addr': '10.45.1.1', + 'addr6': 'dead:beef::1' + } + } + ], + 'security': { + 'k': 'iamatransparentsecretkeystringk', + 'amf': '8000', + 'op': None, + 'opc': 'iamatransparentsecretopcstring' + }, + 'subscribed_rau_tau_timer': 12, + 'subscriber_status': 0 + } + + def test_top_level_migration(self): + new_sub = SchemaUpdater.create_v1_from_v0(self.legacy_sub) + self.assertEqual(new_sub["imsi"], self.legacy_sub["imsi"]) + self.assertEqual(new_sub["subscriber_status"], self.legacy_sub["subscriber_status"]) + self.assertEqual(new_sub["subscribed_rau_tau_timer"], self.legacy_sub["subscribed_rau_tau_timer"]) + self.assertEqual(new_sub["network_access_mode"], self.legacy_sub["network_access_mode"]) + self.assertEqual(new_sub["access_restriction_data"], self.legacy_sub["access_restriction_data"]) + + def test_ambr_migration(self): + new_sub = SchemaUpdater.create_v1_from_v0(self.legacy_sub) + self.assertEqual( + new_sub["ambr"]["uplink"]["value"] * (1000 ** new_sub["ambr"]["uplink"]["unit"]), + self.legacy_sub["ambr"]["uplink"] + ) + + self.assertEqual( + new_sub["ambr"]["downlink"]["value"] * (1000 ** new_sub["ambr"]["downlink"]["unit"]), + self.legacy_sub["ambr"]["downlink"] + ) + + def test_pdn_migration(self): + new_sub = SchemaUpdater.create_v1_from_v0(self.legacy_sub) + self.assertEqual(len(new_sub["slice"]), 1) + self.assertEqual(len(new_sub["slice"][0]["session"]), len(self.legacy_sub["pdn"])) + + session = new_sub["slice"][0]["session"][0] + self.assertEqual( + session["ambr"]["uplink"]["value"] * (1000 ** session["ambr"]["uplink"]["unit"]), + self.legacy_sub["pdn"][0]["ambr"]["uplink"] + ) + + self.assertEqual( + session["ambr"]["downlink"]["value"] * (1000 ** session["ambr"]["downlink"]["unit"]), + self.legacy_sub["pdn"][0]["ambr"]["downlink"] + ) + + self.assertEqual( + session["ue"]["addr"], + self.legacy_sub["pdn"][0]["ue"]["addr"] + ) + + self.assertEqual( + session["ue"]["addr6"], + self.legacy_sub["pdn"][0]["ue"]["addr6"] + )