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 send_abort_session(client, args): abort_session_chan = ServiceRegistry.get_rpc_channel("abort_session_service", ServiceRegistry.LOCAL) abort_session_client = AbortSessionResponderStub(abort_session_chan) asr = AbortSessionRequest( session_id=args.session_id, user_name=args.user_name, ) asa = abort_session_client.AbortSession(asr) print(asa)
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, ))