def process_update(self, stream_name, updates, resync): """ The cloud streams ALL subscribers registered, both active and inactive. Since we don't have a good way of knowing whether a detach succeeds or fails to update the local database correctly, we have to send down all subscribers to keep trying to delete inactive subscribers. TODO we can optimize a bit on the MME side to not detach already detached subscribers. """ logging.info("Processing %d subscriber updates (resync=%s)", len(updates), resync) if resync: # TODO: # - handle database exceptions keys = [] subscribers = [] active_subscriber_ids = [] for update in updates: sub = SubscriberData() sub.ParseFromString(update.value) subscribers.append(sub) keys.append(update.key) if sub.lte.state == LTESubscription.ACTIVE: active_subscriber_ids.append(update.key) old_sub_ids = self._store.list_subscribers() # Only compare active subscribers against the database to decide # what to detach. self.detach_deleted_subscribers(old_sub_ids, active_subscriber_ids) logging.debug("Resync with subscribers: %s", ','.join(keys)) self._store.resync(subscribers) else: # TODO: implement updates pass
def get_subscriber_data(self, subscriber_id): """ Return the auth key for the subscriber. """ db_location = self._db_locations[self._sid2bucket(subscriber_id)] conn = sqlite3.connect(db_location, uri=True) try: with conn: res = conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (subscriber_id, ), ) row = res.fetchone() if not row: raise SubscriberNotFoundError(subscriber_id) except sqlite3.OperationalError: # Print the process holding the lock db_parts = db_location.split(":", 1) if (len(db_parts) == 2) and db_parts[1]: path_str = db_parts[1].split("?") output = subprocess.Popen( ["/usr/bin/fuser", "-uv", path_str[0]], stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) logging.info(output.communicate()) raise SubscriberServerTooBusy(subscriber_id) finally: conn.close() subscriber_data = SubscriberData() subscriber_data.ParseFromString(row[0]) return subscriber_data
def edit_subscriber(self, subscriber_id): """ Context manager to modify the subscriber data. """ db_location = self._db_locations[self._sid2bucket(subscriber_id)] conn = sqlite3.connect(db_location, uri=True) try: with conn: res = conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (subscriber_id, ), ) row = res.fetchone() if not row: raise SubscriberNotFoundError(subscriber_id) subscriber_data = SubscriberData() subscriber_data.ParseFromString(row[0]) yield subscriber_data data_str = subscriber_data.SerializeToString() conn.execute( "UPDATE subscriberdb SET data = ? " "WHERE subscriber_id = ?", (data_str, subscriber_id), ) finally: conn.close()
def resync(self, subscribers): """ Method that should resync the store with the mentioned list of subscribers. The resync leaves the current state of subscribers intact. Args: subscribers - list of subscribers to be in the store. """ with self.conn: # Capture the current state of the subscribers res = self.conn.execute( "SELECT subscriber_id, data FROM subscriberdb") current_state = {} for row in res: sub = SubscriberData() sub.ParseFromString(row[1]) current_state[row[0]] = sub.state # Clear all subscribers self.conn.execute("DELETE FROM subscriberdb") # Add the subscribers with the current state for sub in subscribers: sid = SIDUtils.to_str(sub.sid) if sid in current_state: sub.state.CopyFrom(current_state[sid]) data_str = sub.SerializeToString() self.conn.execute( "INSERT INTO subscriberdb(subscriber_id, data) " "VALUES (?, ?)", (sid, data_str)) self._on_ready.resync(subscribers)
def edit_subscriber(self, subscriber_id): """ Context manager to modify the subscriber data. """ start = time.clock_gettime(time.CLOCK_THREAD_CPUTIME_ID) start2 = time.clock_gettime(time.CLOCK_MONOTONIC) with self.conn: res = self.conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (subscriber_id, ), ) row = res.fetchone() if not row: raise SubscriberNotFoundError(subscriber_id) subscriber_data = SubscriberData() subscriber_data.ParseFromString(row[0]) yield subscriber_data data_str = subscriber_data.SerializeToString() self.conn.execute( "UPDATE subscriberdb SET data = ? " "WHERE subscriber_id = ?", (data_str, subscriber_id), ) end = time.clock_gettime(time.CLOCK_THREAD_CPUTIME_ID) end2 = time.clock_gettime(time.CLOCK_MONOTONIC) logging.warning('LTE edit sub spends: {} ms'.format( (end - start) * 1e3)) logging.warning('LTE edit sub takes: {} ms'.format( (end2 - start2) * 1e3))
def setUp(self): # Create sqlite3 database for testing self._tmpfile = tempfile.TemporaryDirectory() store = SqliteStore(self._tmpfile.name + '/') op = 16 * b'\x11' amf = b'\x80\x00' self._sub_profiles = { 'superfast': SubscriberDB.SubscriptionProfile(max_ul_bit_rate=100000, max_dl_bit_rate=50000) } self._default_sub_profile = SubscriberDB.SubscriptionProfile( max_ul_bit_rate=10000, max_dl_bit_rate=5000) self._processor = processor.Processor(store, self._default_sub_profile, self._sub_profiles, op, amf) # Add some test users (rand, sres, gsm_key) = _dummy_auth_tuple() gsm = GSMSubscription(state=GSMSubscription.ACTIVE, auth_tuples=[rand + sres + gsm_key]) lte_key = 16 * b'\x00' lte = LTESubscription(state=LTESubscription.ACTIVE, auth_key=lte_key) lte_opc = LTESubscription(state=LTESubscription.ACTIVE, auth_key=lte_key, auth_opc=Milenage.generate_opc(lte_key, op)) lte_opc_short = LTESubscription(state=LTESubscription.ACTIVE, auth_key=lte_key, auth_opc=b'\x00') state = SubscriberState(lte_auth_next_seq=1) sub1 = SubscriberData(sid=SIDUtils.to_pb('IMSI11111'), gsm=gsm, lte=lte, state=state, sub_profile='superfast') sub2 = SubscriberData( sid=SIDUtils.to_pb('IMSI22222'), # No auth keys gsm=GSMSubscription(state=GSMSubscription.ACTIVE), lte=LTESubscription(state=LTESubscription.ACTIVE)) sub3 = SubscriberData( sid=SIDUtils.to_pb('IMSI33333')) # No subscribtion sub4 = SubscriberData(sid=SIDUtils.to_pb('IMSI44444'), lte=lte_opc, state=state) sub5 = SubscriberData(sid=SIDUtils.to_pb('IMSI55555'), lte=lte_opc_short, state=state) store.add_subscriber(sub1) store.add_subscriber(sub2) store.add_subscriber(sub3) store.add_subscriber(sub4) store.add_subscriber(sub5)
def get_subscriber_data(self, subscriber_id): """ Method that returns the auth key for the subscriber. """ with self.conn: res = self.conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (subscriber_id, )) row = res.fetchone() if not row: raise SubscriberNotFoundError(subscriber_id) subscriber_data = SubscriberData() subscriber_data.ParseFromString(row[0]) return subscriber_data
def add_subscriber(self, subscriber_data: SubscriberData): """ Add the subscriber to store. """ sid = SIDUtils.to_str(subscriber_data.sid) data_str = subscriber_data.SerializeToString() db_location = self._db_locations[self._sid2bucket(sid)] conn = sqlite3.connect(db_location, uri=True) try: with conn: res = conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (sid, ), ) if res.fetchone(): raise DuplicateSubscriberError(sid) conn.execute( "INSERT INTO subscriberdb(subscriber_id, data) " "VALUES (?, ?)", (sid, data_str), ) finally: conn.close() self._on_ready.add_subscriber(subscriber_data)
def upsert_subscriber(self, subscriber_data: SubscriberData) -> None: """ Check if the given subscriber exists in store. If so, update subscriber data; otherwise, add subscriber. """ sid = SIDUtils.to_str(subscriber_data.sid) data_str = subscriber_data.SerializeToString() db_location = self._db_locations[self._sid2bucket(sid)] conn = sqlite3.connect(db_location, uri=True) try: with conn: res = conn.execute( "SELECT subscriber_id FROM subscriberdb WHERE " "subscriber_id = ?", (sid, ), ) row = res.fetchone() if row is None: conn.execute( "INSERT INTO subscriberdb(subscriber_id, data) " "VALUES (?, ?)", (sid, data_str), ) else: conn.execute( "UPDATE subscriberdb SET data = ? " "WHERE subscriber_id = ?", (data_str, sid), ) finally: conn.close() self._on_ready.upsert_subscriber(subscriber_data)
def add_subscriber(client, args): gsm = GSMSubscription() lte = LTESubscription() state = SubscriberState() if len(args.gsm_auth_tuple) != 0: gsm.state = GSMSubscription.ACTIVE for auth_tuple in args.gsm_auth_tuple: gsm.auth_tuples.append(bytes.fromhex(auth_tuple)) if args.lte_auth_key is not None: lte.state = LTESubscription.ACTIVE lte.auth_key = bytes.fromhex(args.lte_auth_key) if args.lte_auth_next_seq is not None: state.lte_auth_next_seq = args.lte_auth_next_seq if args.lte_auth_opc is not None: lte.auth_opc = bytes.fromhex(args.lte_auth_opc) data = SubscriberData( sid=SIDUtils.to_pb(args.sid), gsm=gsm, lte=lte, state=state, ) client.AddSubscriber(data)
def add_incomplete_sub(cls, sid: str): sub_db_sid = SIDUtils.to_pb(sid) lte = LTESubscription() lte.state = LTESubscription.ACTIVE state = SubscriberState() state.lte_auth_next_seq = 1 subs_data = SubscriberData(sid=sub_db_sid, lte=lte, state=state) cls.subs[str(sub_db_sid)] = subs_data
def edit_subscriber(self, subscriber_id): """ Context manager to modify the subscriber data. """ with self.conn: res = self.conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (subscriber_id, )) row = res.fetchone() if not row: raise SubscriberNotFoundError(subscriber_id) subscriber_data = SubscriberData() subscriber_data.ParseFromString(row[0]) yield subscriber_data data_str = subscriber_data.SerializeToString() self.conn.execute( "UPDATE subscriberdb SET data = ? " "WHERE subscriber_id = ?", (data_str, subscriber_id))
def resync(self, subscribers): """ Method that should resync the store with the mentioned list of subscribers. The resync leaves the current state of subscribers intact. Args: subscribers - list of subscribers to be in the store. """ bucket_subs = defaultdict(list) for sub in subscribers: sid = SIDUtils.to_str(sub.sid) bucket_subs[self._sid2bucket(sid)].append(sub) for i, db_location in enumerate(self._db_locations): conn = sqlite3.connect(db_location, uri=True) try: with conn: # Capture the current state of the subscribers res = conn.execute( "SELECT subscriber_id, data FROM subscriberdb", ) current_state = {} for row in res: sub = SubscriberData() sub.ParseFromString(row[1]) current_state[row[0]] = sub.state # Clear all subscribers conn.execute("DELETE FROM subscriberdb") # Add the subscribers with the current state for sub in bucket_subs[i]: sid = SIDUtils.to_str(sub.sid) if sid in current_state: sub.state.CopyFrom(current_state[sid]) data_str = sub.SerializeToString() conn.execute( "INSERT INTO subscriberdb(subscriber_id, data) " "VALUES (?, ?)", (sid, data_str), ) finally: conn.close() self._on_ready.resync(subscribers)
def get_subscriber_data(self, subscriber_id): """ Method that returns the auth key for the subscriber. """ db_location = self._db_locations[self._sid2bucket(subscriber_id)] conn = sqlite3.connect(db_location, uri=True) try: with conn: res = conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (subscriber_id, )) row = res.fetchone() if not row: raise SubscriberNotFoundError(subscriber_id) finally: conn.close() subscriber_data = SubscriberData() subscriber_data.ParseFromString(row[0]) return subscriber_data
def _build_add_subs_data(num_subs: int, input_file: str): add_subs_reqs = [] for i in range(1, num_subs): sid = SubscriberID(id=str(i).zfill(15)) config = Non3GPPUserProfile( apn_config=[APNConfiguration(service_selection=TEST_APN)], ) data = SubscriberData(sid=sid, non_3gpp=config) add_sub_req_dict = json_format.MessageToDict(data) add_subs_reqs.append(add_sub_req_dict) with open(input_file, 'w') as file: json.dump(add_subs_reqs, file, separators=(',', ':'))
def get_all_subscribers(self): return [ SubscriberData( sid=SubscriberID( id="IMSI111", ), ), SubscriberData( sid=SubscriberID( id="IMSI222", ), ), SubscriberData( sid=SubscriberID( id="IMSI333", ), ), SubscriberData( sid=SubscriberID( id="IMSI444", ), ), SubscriberData( sid=SubscriberID( id="IMSI555", ), ), SubscriberData( sid=SubscriberID( id="IMSI666", ), ), ]
def add_srsue(self): # Add for srsUE, add the hardcoded SubscriberData into the store when receiving request imsi = 'IMSI' + SRSUE_IMSI lte = LTESubscription() lte.state = LTESubscription.ACTIVE lte.auth_key = bytes.fromhex(SRSUE_KEY) #lte.auth_opc = bytes.fromhex(SRSUE_OPC) state = SubscriberState() state.lte_auth_next_seq = 1 sub_data = SubscriberData(sid=SIDUtils.to_pb(imsi), lte=lte, state=state) self.add_subscriber(sub_data)
def _build_update_subs_data(num_subs: int, input_file: str): active_sids = _load_subs(num_subs) update_subs_reqs = [] for sid in active_sids: config = Non3GPPUserProfile( apn_config=[APNConfiguration(service_selection=TEST_APN_UPDATE)], ) data = SubscriberData(sid=sid, non_3gpp=config) update = SubscriberUpdate(data=data) update_sub_req_dict = json_format.MessageToDict(update) update_subs_reqs.append(update_sub_req_dict) with open(input_file, 'w') as file: json.dump(update_subs_reqs, file, separators=(',', ':'))
def process_update(self, stream_name, updates, resync): logging.info("Processing %d subscriber updates (resync=%s)", len(updates), resync) if resync: # TODO: # - handle database exceptions keys = [] subscribers = [] for update in updates: sub = SubscriberData() sub.ParseFromString(update.value) subscribers.append(sub) keys.append(update.key) old_sub_ids = self._store.list_subscribers() self.detach_deleted_subscribers(old_sub_ids, keys) logging.debug("Resync with subscribers: %s", ','.join(keys)) self._store.resync(subscribers) else: # TODO: implement updates pass
def add_sub(cls, sid: str, apn: str, ip: str, vlan: str = None, gw_ip=None, gw_mac=None): sub_db_sid = SIDUtils.to_pb(sid) lte = LTESubscription() lte.state = LTESubscription.ACTIVE state = SubscriberState() state.lte_auth_next_seq = 1 non_3gpp = Non3GPPUserProfile() subs_data = SubscriberData(sid=sub_db_sid, lte=lte, state=state, non_3gpp=non_3gpp) cls.subs[str(sub_db_sid)] = subs_data cls.add_sub_ip(sid, apn, ip, vlan, gw_ip, gw_mac)
def _load_subs(num_subs: int) -> List[SubscriberID]: client = SubscriberDBStub( ServiceRegistry.get_rpc_channel('subscriberdb', ServiceRegistry.LOCAL)) sids = [] for i in range(1, num_subs): sid = SubscriberID(id=str(i).zfill(15)) config = Non3GPPUserProfile( apn_config=[APNConfiguration(service_selection="magma.ipv4")]) data = SubscriberData(sid=sid, non_3gpp=config) client.AddSubscriber(data) sids.append(sid) return sids
def test_upsert_subscriber(self): """ Test if subscriber upsertion triggers ready """ self.assertEqual(self._store._on_ready.event.is_set(), False) self._store.upsert_subscriber( SubscriberData(sid=SIDUtils.to_pb('IMSI1111')), ) async def defer(): await self._store.on_ready() self.loop.run_until_complete(defer()) self.assertEqual(self._store._on_ready.event.is_set(), True)
def _load_subs(num_subs: int) -> List[SubscriberID]: client = SubscriberDBStub( ServiceRegistry.get_rpc_channel( SUBSCRIBERDB_SERVICE_NAME, ServiceRegistry.LOCAL, ), ) sids = [] for i in range(1, num_subs): sid = SubscriberID(id=str(i).zfill(15)) config = Non3GPPUserProfile( apn_config=[APNConfiguration(service_selection=TEST_APN)], ) data = SubscriberData(sid=sid, non_3gpp=config) client.AddSubscriber(data) sids.append(sid) return sids
def _get_subscriberdb_data(sid): """ Get subscriber data in protobuf format. Args: sid (str): string representation of the subscriber id Returns: subscriber_data (protos.subscriberdb_pb2.SubscriberData): full subscriber information for :sid: in protobuf format. """ sub_db_sid = SIDUtils.to_pb(sid) lte = LTESubscription() lte.state = LTESubscription.ACTIVE lte.auth_key = bytes.fromhex(KEY) state = SubscriberState() state.lte_auth_next_seq = 1 return SubscriberData(sid=sub_db_sid, lte=lte, state=state)
def add_subscriber(client, args): gsm = GSMSubscription() lte = LTESubscription() state = SubscriberState() sub_network = CoreNetworkType() if len(args.gsm_auth_tuple) != 0: gsm.state = GSMSubscription.ACTIVE for auth_tuple in args.gsm_auth_tuple: gsm.auth_tuples.append(bytes.fromhex(auth_tuple)) if args.lte_auth_key is not None: lte.state = LTESubscription.ACTIVE lte.auth_key = bytes.fromhex(args.lte_auth_key) if args.lte_auth_next_seq is not None: state.lte_auth_next_seq = args.lte_auth_next_seq if args.lte_auth_opc is not None: lte.auth_opc = bytes.fromhex(args.lte_auth_opc) if args.forbidden_network_types is not None: if (len(args.forbidden_network_types.split(",")) > 2): print("Forbidden Core Network Types are NT_5GC, NT_EPC") return for n in args.forbidden_network_types.split(","): if n == "NT_5GC": sub_network.forbidden_network_types.extend( [CoreNetworkType.NT_5GC]) elif n == "NT_EPC": sub_network.forbidden_network_types.extend( [CoreNetworkType.NT_EPC]) else: print( "Invalid Network type, Forbidden Core Network Types are NT_5GC, NT_EPC" ) return data = SubscriberData( sid=SIDUtils.to_pb(args.sid), gsm=gsm, lte=lte, state=state, sub_network=sub_network, ) client.AddSubscriber(data)
def ListSubscribers(self, request: ListSubscribersRequest, context) -> ListSubscribersResponse: # noqa: N802 """ List subscribers is a mock to trigger various test cases Args: request: ListSubscribersRequest context: request context Raises: RpcError: If page size is 1 Returns: ListSubscribersResponse """ # Add in logic to allow error handling testing root_digest = Digest(md5_base64_digest="") leaf_digests = [] if request.page_size == 1: raise grpc.RpcError("Test Exception") if request.page_token == "": next_page_token = "aaa" # noqa: S105 subscribers = [ SubscriberData(sid=SubscriberID(id="IMSI111")), SubscriberData(sid=SubscriberID(id="IMSI222")), ] root_digest = Digest(md5_base64_digest="root_digest_apple") leaf_digests = [ LeafDigest( id='IMSI11111', digest=Digest(md5_base64_digest="leaf_digests_apple"), ), ] elif request.page_token == "aaa": next_page_token = "bbb" # noqa: S105 subscribers = [ SubscriberData(sid=SubscriberID(id="IMSI333")), SubscriberData(sid=SubscriberID(id="IMSI444")), ] else: next_page_token = "" # noqa: S105 subscribers = [ SubscriberData(sid=SubscriberID(id="IMSI555")), SubscriberData(sid=SubscriberID(id="IMSI666")), ] return ListSubscribersResponse( subscribers=subscribers, next_page_token=next_page_token, digests=DigestTree( root_digest=root_digest, leaf_digests=leaf_digests, ), )
def _process_changeset( self, changeset: Optional[Changeset]) -> ProcessedChangeset: if changeset is None: return ProcessedChangeset(to_renew=[], deleted=[]) to_renew, deleted = [], [] if changeset.deleted is not None: deleted = changeset.deleted if changeset.to_renew is not None: for any_val in changeset.to_renew: data = SubscriberData() ok = any_val.Unpack(data) if not ok: raise ValueError( 'Cannot unpack Any type into message: %s' % data, ) to_renew.append(data) return ProcessedChangeset(to_renew=to_renew, deleted=deleted)
async def test(): # noqa: WPS430 get_grpc_mock.return_value = self.channel subscribers = (await self.subscriberdb_cloud_client._get_subscribers()) expected_subscribers = [ SubscriberData(sid=SubscriberID(id="IMSI111", ), ), SubscriberData(sid=SubscriberID(id="IMSI222", ), ), SubscriberData(sid=SubscriberID(id="IMSI333", ), ), SubscriberData(sid=SubscriberID(id="IMSI444", ), ), SubscriberData(sid=SubscriberID(id="IMSI555", ), ), SubscriberData(sid=SubscriberID(id="IMSI666", ), ), ] self.assertEqual(expected_subscribers, subscribers)
def ListSubscribers(self, request, context): # noqa: N802 """ List subscribers is a mock to trigger various test cases Args: request: ListSubscribersRequest context: request context Raises: RpcError: If page size is 1 Returns: ListSubscribersResponse """ # Add in logic to allow error handling testing flat_digest = Digest(md5_base64_digest="") per_sub_digests = [] if request.page_size == 1: raise grpc.RpcError("Test Exception") if request.page_token == "": next_page_token = "aaa" # noqa: S105 subscribers = [ SubscriberData(sid=SubscriberID(id="IMSI111")), SubscriberData(sid=SubscriberID(id="IMSI222")), ] flat_digest = Digest(md5_base64_digest="flat_digest_apple") per_sub_digests = [ SubscriberDigestWithID( sid=SIDUtils.to_pb("IMSI11111"), digest=Digest(md5_base64_digest="per_sub_digests_apple"), ), ] elif request.page_token == "aaa": next_page_token = "bbb" # noqa: S105 subscribers = [ SubscriberData(sid=SubscriberID(id="IMSI333")), SubscriberData(sid=SubscriberID(id="IMSI444")), ] else: next_page_token = "" # noqa: S105 subscribers = [ SubscriberData(sid=SubscriberID(id="IMSI555")), SubscriberData(sid=SubscriberID(id="IMSI666")), ] return ListSubscribersResponse( subscribers=subscribers, next_page_token=next_page_token, flat_digest=flat_digest, per_sub_digests=per_sub_digests, )
def test_update_subscriber(self): """ Test if UpdateSubscriber rpc call works """ sid = SIDUtils.to_pb('IMSI1') data = SubscriberData(sid=sid) # Add subscriber self._stub.AddSubscriber(data) sub = self._stub.GetSubscriberData(sid) self.assertEqual(sub.lte.auth_key, b'') self.assertEqual(sub.state.lte_auth_next_seq, 0) # Update subscriber update = SubscriberUpdate() update.data.sid.CopyFrom(sid) update.data.lte.auth_key = b'\xab\xcd' update.data.state.lte_auth_next_seq = 1 update.mask.paths.append('lte.auth_key') # only auth_key self._stub.UpdateSubscriber(update) sub = self._stub.GetSubscriberData(sid) self.assertEqual(sub.state.lte_auth_next_seq, 0) # no change self.assertEqual(sub.lte.auth_key, b'\xab\xcd') update.data.state.lte_auth_next_seq = 1 update.mask.paths.append('state.lte_auth_next_seq') self._stub.UpdateSubscriber(update) sub = self._stub.GetSubscriberData(sid) self.assertEqual(sub.state.lte_auth_next_seq, 1) # Delete the subscriber self._stub.DeleteSubscriber(sid) with self.assertRaises(grpc.RpcError) as err: self._stub.UpdateSubscriber(update) self.assertEqual(err.exception.code(), grpc.StatusCode.NOT_FOUND)