def update_record(imsi: str, ip_addr: str) -> None:
"""
Make RPC call to 'UpdateRecord' method of local directoryD service
"""
try:
chan = ServiceRegistry.get_rpc_channel(DIRECTORYD_SERVICE_NAME,
ServiceRegistry.LOCAL)
except ValueError:
logging.error('Cant get RPC channel to %s', DIRECTORYD_SERVICE_NAME)
return
client = GatewayDirectoryServiceStub(chan)
if not imsi.startswith("IMSI"):
imsi = "IMSI" + imsi
try:
# Location will be filled in by directory service
req = UpdateRecordRequest(id=imsi, location="hwid")
req.fields[IPV4_ADDR_KEY] = ip_addr
client.UpdateRecord(req, DEFAULT_GRPC_TIMEOUT)
except grpc.RpcError as err:
logging.error(
"UpdateRecordRequest error for id: %s, ipv4_addr: %s! [%s] %s",
imsi,
ip_addr,
err.code(),
err.details())
def get_record(imsi: str, field: str) -> str:
"""
Make RPC call to 'GetDirectoryField' method of local directoryD service
"""
try:
chan = ServiceRegistry.get_rpc_channel(DIRECTORYD_SERVICE_NAME,
ServiceRegistry.LOCAL)
except ValueError:
logging.error('Cant get RPC channel to %s', DIRECTORYD_SERVICE_NAME)
return
client = GatewayDirectoryServiceStub(chan)
if not imsi.startswith("IMSI"):
imsi = "IMSI" + imsi
try:
# Location will be filled in by directory service
req = GetDirectoryFieldRequest(id=imsi, field_key=field)
res = client.GetDirectoryField(req, DEFAULT_GRPC_TIMEOUT)
if res.value is not None:
return res.value
except grpc.RpcError as err:
logging.error(
"GetDirectoryFieldRequest error for id: %s! [%s] %s",
imsi,
err.code(),
err.details())
return None
def get_all_records(retries: int = 3, sleep_time: float = 0.1) -> [dict]:
"""
Make RPC call to 'GetAllDirectoryRecords' method of local directoryD service
"""
try:
chan = ServiceRegistry.get_rpc_channel(
DIRECTORYD_SERVICE_NAME,
ServiceRegistry.LOCAL,
)
except ValueError:
logging.error('Cant get RPC channel to %s', DIRECTORYD_SERVICE_NAME)
return
client = GatewayDirectoryServiceStub(chan)
for _ in range(0, retries):
try:
res = client.GetAllDirectoryRecords(Void(), DEFAULT_GRPC_TIMEOUT)
if res.records is not None:
return res.records
hub.sleep(sleep_time)
except grpc.RpcError as err:
logging.error(
"GetAllDirectoryRecords error! [%s] %s",
err.code(),
err.details(),
)
return []
def __init__(self):
"""
Initialize sessionManager util.
"""
self._session_stub = SessionProxyResponderStub(
get_rpc_channel("sessiond"))
self._directorydstub = GatewayDirectoryServiceStub(
get_rpc_channel("directoryd"))
def _cleanup_subs():
"""Clear directory records"""
client = GatewayDirectoryServiceStub(
ServiceRegistry.get_rpc_channel(
DIRECTORYD_SERVICE_NAME,
ServiceRegistry.LOCAL,
), )
for record in client.GetAllDirectoryRecords(Void()).records:
sid = DeleteRecordRequest(id=record.id, )
client.DeleteRecord(sid)
def setUp(self):
# Bind the rpc server to a free port
thread_pool = futures.ThreadPoolExecutor(max_workers=10)
self._rpc_server = grpc.server(thread_pool)
port = self._rpc_server.add_insecure_port('0.0.0.0:0')
# Add the servicer
self._servicer = GatewayDirectoryServiceRpcServicer()
self._servicer.add_to_server(self._rpc_server)
self._rpc_server.start()
# Create a rpc stub
channel = grpc.insecure_channel('0.0.0.0:{}'.format(port))
self._stub = GatewayDirectoryServiceStub(channel)
def main():
""" main() for smsd """
service = MagmaService('smsd', None)
# Optionally pipe errors to Sentry
sentry_init(service_name=service.name)
directoryd_chan = ServiceRegistry.get_rpc_channel(
'directoryd',
ServiceRegistry.LOCAL,
)
mme_chan = ServiceRegistry.get_rpc_channel(
'sms_mme_service',
ServiceRegistry.LOCAL,
)
smsd_chan = ServiceRegistry.get_rpc_channel('smsd', ServiceRegistry.CLOUD)
# Add all servicers to the server
smsd_relay = SmsRelay(
service.loop,
GatewayDirectoryServiceStub(directoryd_chan),
SMSOrc8rGatewayServiceStub(mme_chan),
SmsDStub(smsd_chan),
)
smsd_relay.add_to_server(service.rpc_server)
smsd_relay.start()
# Run the service loop
service.run()
# Cleanup the service
service.close()
def setUp(self):
# Bind the rpc server to a free port
thread_pool = futures.ThreadPoolExecutor(max_workers=10)
self._rpc_server = grpc.server(thread_pool)
port = self._rpc_server.add_insecure_port('0.0.0.0:0')
# mock the get_default_client function used to return the same
# fakeredis object
func_mock = \
mock.MagicMock(return_value=fakeredis.FakeStrictRedis())
with mock.patch('magma.directoryd.rpc_servicer.get_default_client',
func_mock):
# Add the servicer
self._servicer = GatewayDirectoryServiceRpcServicer(False)
self._servicer.add_to_server(self._rpc_server)
self._rpc_server.start()
# Create a rpc stub
channel = grpc.insecure_channel('0.0.0.0:{}'.format(port))
self._stub = GatewayDirectoryServiceStub(channel)
def _load_subs(num_subs: int) -> List[DirectoryRecord]:
"""Load directory records"""
client = GatewayDirectoryServiceStub(
ServiceRegistry.get_rpc_channel(
DIRECTORYD_SERVICE_NAME,
ServiceRegistry.LOCAL,
), )
sids = []
for i in range(num_subs):
mac_addr = (str(i) * 2 + ":") * 5 + (str(i) * 2)
ipv4_addr = str(i) * 3 + "." + str(i) * 3 + "." + str(
i) * 3 + "." + str(i) * 3
fields = {"mac-addr": mac_addr, "ipv4_addr": ipv4_addr}
sid = UpdateRecordRequest(
fields=fields,
id=str(i).zfill(15),
location=str(i).zfill(15),
)
client.UpdateRecord(sid)
sids.append(sid)
return sids
class SessionManagerUtil(object):
"""
Helper class to communicate with session manager for the tests.
"""
def __init__(self):
"""
Initialize sessionManager util.
"""
self._session_stub = SessionProxyResponderStub(
get_rpc_channel("sessiond"))
self._abort_session_stub = AbortSessionResponderStub(
get_rpc_channel("abort_session_service"))
self._directorydstub = GatewayDirectoryServiceStub(
get_rpc_channel("directoryd"))
def get_flow_match(self, flow_list, flow_match_list):
"""
Populates flow match list
"""
for flow in flow_list:
flow_direction = flow["direction"]
ip_protocol = flow["ip_proto"]
if ip_protocol == FlowMatch.IPPROTO_TCP:
udp_src_port = 0
udp_dst_port = 0
tcp_src_port = (int(flow["tcp_src_port"])
if "tcp_src_port" in flow else 0)
tcp_dst_port = (int(flow["tcp_dst_port"])
if "tcp_dst_port" in flow else 0)
elif ip_protocol == FlowMatch.IPPROTO_UDP:
tcp_src_port = 0
tcp_dst_port = 0
udp_src_port = (int(flow["udp_src_port"])
if "udp_src_port" in flow else 0)
udp_dst_port = (int(flow["udp_dst_port"])
if "udp_dst_port" in flow else 0)
else:
udp_src_port = 0
udp_dst_port = 0
tcp_src_port = 0
tcp_dst_port = 0
src_addr = None
if flow.get("ipv4_src", None):
src_addr = IPAddress(
version=IPAddress.IPV4,
address=flow.get("ipv4_src").encode('utf-8'))
elif flow.get("ipv6_src", None):
src_addr = IPAddress(
version=IPAddress.IPV6,
address=flow.get("ipv6_src").encode('utf-8'))
dst_addr = None
if flow.get("ipv4_dst", None):
dst_addr = IPAddress(
version=IPAddress.IPV4,
address=flow.get("ipv4_dst").encode('utf-8'))
elif flow.get("ipv6_dst", None):
dst_addr = IPAddress(
version=IPAddress.IPV6,
address=flow.get("ipv6_dst").encode('utf-8'))
flow_match_list.append(
FlowDescription(
match=FlowMatch(
ip_dst=dst_addr,
ip_src=src_addr,
tcp_src=tcp_src_port,
tcp_dst=tcp_dst_port,
udp_src=udp_src_port,
udp_dst=udp_dst_port,
ip_proto=ip_protocol,
direction=flow_direction,
),
action=FlowDescription.PERMIT,
))
def create_ReAuthRequest(self, imsi, policy_id, flow_list, qos):
"""
Sends Policy RAR message to session manager
"""
print("Sending Policy RAR message to session manager")
flow_match_list = []
res = None
self.get_flow_match(flow_list, flow_match_list)
policy_qos = FlowQos(
qci=qos["qci"],
max_req_bw_ul=qos["max_req_bw_ul"],
max_req_bw_dl=qos["max_req_bw_dl"],
gbr_ul=qos["gbr_ul"],
gbr_dl=qos["gbr_dl"],
arp=QosArp(
priority_level=qos["arp_prio"],
pre_capability=qos["pre_cap"],
pre_vulnerability=qos["pre_vul"],
),
)
policy_rule = PolicyRule(
id=policy_id,
priority=qos["priority"],
flow_list=flow_match_list,
tracking_type=PolicyRule.NO_TRACKING,
rating_group=1,
monitoring_key=None,
qos=policy_qos,
)
qos = QoSInformation(qci=qos["qci"])
# Get sessionid
req = GetDirectoryFieldRequest(id=imsi, field_key="session_id")
try:
res = self._directorydstub.GetDirectoryField(
req, DEFAULT_GRPC_TIMEOUT)
except grpc.RpcError as err:
logging.error(
"GetDirectoryFieldRequest error for id: %s! [%s] %s",
imsi,
err.code(),
err.details(),
)
self._session_stub.PolicyReAuth(
PolicyReAuthRequest(
session_id=res.value,
imsi=imsi,
rules_to_remove=[],
rules_to_install=[],
dynamic_rules_to_install=[
DynamicRuleInstall(policy_rule=policy_rule)
],
event_triggers=[],
revalidation_time=None,
usage_monitoring_credits=[],
qos_info=qos,
))
def create_AbortSessionRequest(self, imsi: str) -> AbortSessionResult:
# Get SessionID
req = GetDirectoryFieldRequest(id=imsi, field_key="session_id")
try:
res = self._directorydstub.GetDirectoryField(
req, DEFAULT_GRPC_TIMEOUT)
except grpc.RpcError as err:
logging.error(
"GetDirectoryFieldRequest error for id: %s! [%s] %s",
imsi,
err.code(),
err.details(),
)
return self._abort_session_stub.AbortSession(
AbortSessionRequest(
session_id=res.value,
user_name=imsi,
))
class DirectorydRpcServiceTests(TestCase):
@mock.patch("redis.Redis", MockRedis)
def setUp(self):
# Bind the rpc server to a free port
thread_pool = futures.ThreadPoolExecutor(max_workers=10)
self._rpc_server = grpc.server(thread_pool)
port = self._rpc_server.add_insecure_port('0.0.0.0:0')
# Add the servicer
self._servicer = GatewayDirectoryServiceRpcServicer()
self._servicer.add_to_server(self._rpc_server)
self._rpc_server.start()
# Create a rpc stub
channel = grpc.insecure_channel('0.0.0.0:{}'.format(port))
self._stub = GatewayDirectoryServiceStub(channel)
@mock.patch("redis.Redis", MockRedis)
def tearDown(self):
self._rpc_server.stop(None)
@mock.patch("redis.Redis", MockRedis)
@mock.patch('snowflake.snowflake', get_mock_snowflake)
def test_update_record(self):
self._servicer._redis_dict.clear()
req = UpdateRecordRequest()
req.id = "IMSI555"
self._stub.UpdateRecord(req)
actual_record = self._servicer._redis_dict[req.id]
self.assertEqual(actual_record.location_history, ['aaa-bbb'])
self.assertEqual(actual_record.identifiers, {})
req.fields["mac_addr"] = "aa:aa:bb:bb:cc:cc"
req.fields["ipv4_addr"] = "192.168.172.12"
self._stub.UpdateRecord(req)
actual_record2 = self._servicer._redis_dict[req.id]
self.assertEqual(actual_record2.location_history, ["aaa-bbb"])
self.assertEqual(actual_record2.identifiers['mac_addr'],
"aa:aa:bb:bb:cc:cc")
self.assertEqual(actual_record2.identifiers['ipv4_addr'],
"192.168.172.12")
@mock.patch("redis.Redis", MockRedis)
@mock.patch('snowflake.snowflake', get_mock_snowflake)
def test_update_record_bad_location(self):
self._servicer._redis_dict.clear()
req = UpdateRecordRequest()
req.id = "IMSI556"
req.location = "bbb-ccc"
self._stub.UpdateRecord(req)
actual_record = self._servicer._redis_dict[req.id]
self.assertEqual(actual_record.location_history, ['aaa-bbb'])
self.assertEqual(actual_record.identifiers, {})
@mock.patch("redis.Redis", MockRedis)
@mock.patch('snowflake.snowflake', get_mock_snowflake)
def test_delete_record(self):
self._servicer._redis_dict.clear()
req = UpdateRecordRequest()
req.id = "IMSI557"
self._stub.UpdateRecord(req)
self.assertTrue(req.id in self._servicer._redis_dict)
del_req = DeleteRecordRequest()
del_req.id = "IMSI557"
self._stub.DeleteRecord(del_req)
self.assertFalse(req.id in self._servicer._redis_dict)
with self.assertRaises(grpc.RpcError) as err:
self._stub.DeleteRecord(del_req)
self.assertEqual(err.exception.code(), grpc.StatusCode.NOT_FOUND)
@mock.patch("redis.Redis", MockRedis)
@mock.patch('snowflake.snowflake', get_mock_snowflake)
def test_get_field(self):
self._servicer._redis_dict.clear()
req = UpdateRecordRequest()
req.id = "IMSI557"
req.fields["mac_addr"] = "aa:bb:aa:bb:aa:bb"
self._stub.UpdateRecord(req)
self.assertTrue(req.id in self._servicer._redis_dict)
get_req = GetDirectoryFieldRequest()
get_req.id = "IMSI557"
get_req.field_key = "mac_addr"
ret = self._stub.GetDirectoryField(get_req)
self.assertEqual("aa:bb:aa:bb:aa:bb", ret.value)
with self.assertRaises(grpc.RpcError) as err:
get_req.field_key = "ipv4_addr"
self._stub.GetDirectoryField(get_req)
self.assertEqual(err.exception.code(), grpc.StatusCode.NOT_FOUND)
@mock.patch("redis.Redis", MockRedis)
@mock.patch('snowflake.snowflake', get_mock_snowflake)
def test_get_all(self):
self._servicer._redis_dict.clear()
req = UpdateRecordRequest()
req.id = "IMSI557"
req.fields["mac_addr"] = "aa:bb:aa:bb:aa:bb"
self._stub.UpdateRecord(req)
self.assertTrue(req.id in self._servicer._redis_dict)
req2 = UpdateRecordRequest()
req2.id = "IMSI556"
req2.fields["ipv4_addr"] = "192.168.127.11"
self._stub.UpdateRecord(req2)
self.assertTrue(req2.id in self._servicer._redis_dict)
void_req = Void()
ret = self._stub.GetAllDirectoryRecords(void_req)
self.assertEqual(2, len(ret.records))
for record in ret.records:
if record.id == "IMSI556":
self.assertEqual(record.fields["ipv4_addr"], "192.168.127.11")
elif record.id == "IMSI557":
self.assertEqual(record.fields["mac_addr"],
"aa:bb:aa:bb:aa:bb")
else:
raise AssertionError()