signaling_filter.py

from __future__ import annotations
from collections.abc import Iterator
from dataclasses import dataclass
from typing import List
import pandas as pd
import argparse,os,fnmatch

@dataclass
class Responses:
    messages = []
    responses={}
    auth_steps={}

    def __init__(self):
        self._populate_auth_steps()
        self._populate_responses()
    
    
    def _populate_auth_steps(self):
        self.auth_steps['InitialUeMessage']='POST-nausf-authv1ue-authentications'
        self.auth_steps['POST-nausf-authv1ue-authentications']='POST-nudm-ueauv1generate-auth-data'
        self.auth_steps['POST-nudm-ueauv1generate-auth-data']='GET-nudr-drv1authentication-subscription'
        self.auth_steps['GET-nudr-drv1authentication-subscription']='AuthenticationRequest'
        self.auth_steps['PATCH-nudr-drv1authentication-subscription']='AuthenticationRequest'
        self.auth_steps['AuthenticationRequest']='AuthenticationResponse'
        self.auth_steps['AuthenticationResponse']='PUT-nausf-authv15g-aka-confirmation'
        self.auth_steps['PUT-nausf-authv15g-aka-confirmation']='POST-nudm-ueauv1auth-events'
        self.auth_steps['POST-nudm-ueauv1auth-events']='PUT-nudr-drv1authentication-status'
        self.auth_steps['PUT-nudr-drv1authentication-status']='SecurityModeCommand'
        self.auth_steps['SecurityModeCommand']='PUT-nudm-uecmv1amf-3gpp-access'
        self.auth_steps['PUT-nudm-uecmv1amf-3gpp-access']='PUT-nudr-drv1amf-3gpp-access'
        self.auth_steps['PUT-nudr-drv1amf-3gpp-access']='GET-nudm-sdmv2am-data'
        self.auth_steps['GET-nudm-sdmv2am-data']='GET-nudr-drv1am-data'
        self.auth_steps['GET-nudr-drv1am-data']='GET-nudr-drv1smf-selection-subscription-data'
        self.auth_steps['GET-nudr-drv1smf-selection-subscription-data']='GET-nudm-sdmv2smf-select-data'
        self.auth_steps['GET-nudm-sdmv2smf-select-data']='GET-nudm-sdmv2ue-context-in-smf-data'
        self.auth_steps['GET-nudm-sdmv2ue-context-in-smf-data']='POST-nudm-sdmv2sdm-subscriptions'
        self.auth_steps['POST-nudm-sdmv2sdm-subscriptions']='POST-npcf-am-policy-controlv1policies'
        self.auth_steps['POST-npcf-am-policy-controlv1policies']='GET-nudr-drv1am-data'
        self.auth_steps['GET-nudr-drv1am-data']='InitialContextSetup'
        self.auth_steps['InitialContextSetup']='POST-nsmf-pdusessionv1sm-contexts'
        self.auth_steps['POST-nsmf-pdusessionv1sm-contexts']='GET-nudm-sdmv2sm-data'
        self.auth_steps['GET-nudm-sdmv2sm-data']='GET-nudr-drv1sm-data'
        self.auth_steps['GET-nudr-drv1sm-data']='POST-npcf-smpolicycontrolv1sm-policies'
        self.auth_steps['POST-npcf-smpolicycontrolv1sm-policies']='GET-nudr-drv1sm-data'
        self.auth_steps['GET-nudr-drv1sm-data']='POST-nbsf-managementv1pcfBindings'
        self.auth_steps['POST-nbsf-managementv1pcfBindings']='SessionEstablishmentRequest'
        self.auth_steps['SessionEstablishmentRequest']='SessionEstablishmentResponse'
        self.auth_steps['SessionEstablishmentResponse']='POST-namf-commv1n1-n2-messages'
        self.auth_steps['POST-namf-commv1n1-n2-messages']='PDUSessionResourceSetup_Req'
        self.auth_steps['PDUSessionResourceSetup_Req']='PDUSessionResourceSetup_Rsp'
        self.auth_steps['PDUSessionResourceSetup_Rsp']='end_placeholder'
        #self.auth_steps['PDUSessionResourceSetup_Req']='PUT-nudm-uecmv1'
        #self.auth_steps['PUT-nudm-uecmv1']='PUT-nudr-drv1'
        #self.auth_steps['PUT-nudr-drv1']='Done'


    def _populate_responses(self):
        self.responses['InitialUeMessage']=['POST-nausf-authv1ue-authentications']
        self.responses['AuthenticationRequest']=['AuthenticationResponse', 'AuthenticationFailure']
        self.responses['AuthenticationFailure']=['POST-nausf-authv1ue-authentications']
        self.responses['POST-nausf-authv1ue-authentications']=['PATCH-nudr-drv1authentication-subscription']
        self.responses['AuthenticationResponse']=['PUT-nausf-authv15g-aka-confirmation']
        self.responses['SecurityModeCommand']=['UplinkNASTransport']
        self.responses['RegistrationReject']=['UplinkNASTransport', 'RegistrationRequest']
        self.responses['RegistrationRequest']=['POST-nausf-authv1ue-authentications']
        self.responses['InitialContextSetup']=['InitialContextSetup']
        self.responses['PDUSessionResourceSetup_Req']=['PDUSessionResourceSetup_Rsp']
        self.responses['PDUSessionResourceSetup_Rsp']=['POST-nsmf-pdusessionv1modify']
        self.responses['POST-nsmf-pdusessionv1modify'] = ['SessionDeletionRequest']
        self.responses['SessionDeletionRequest'] = ['SessionDeletionResponse']
        self.responses['SessionDeletionResponse'] = ['DELETE-nudr-drv1']
        self.responses['DELETE-nudr-drv1'] = ['DELETE-nudm-uecmv1']
        self.responses['DELETE-nudm-uecmv1'] = ['POST-nsmf-pdusessionv1modify']
        self.responses['SessionEstablishmentRequest']=['SessionEstablishmentResponse']
        self.responses['SessionEstablishmentResponse']=['POST-namf-commv1n1-n2-messages','POST-nsmf-pdusessionv1modify']
        self.responses['POST-nsmf-pdusessionv1modify']=['SessionModificationRequest']
        self.responses['SessionModificationRequest']=['SessionModificationResponse']
        self.responses['SessionModificationResponse']=['PUT-nudm-uecmv1']
        self.responses['UplinkNASTransport']=['UplinkNASTransport','DownlinkNASTransport']
        self.responses['DownlinkNASTransport']=['UplinkNASTransport','DownlinkNASTransport']
        self.responses['HTTP']=['200','201','204']



@dataclass
class Message:
    src: str
    dst: str
    in_time: float
    out_time : float
    info: str
    id: str
    ue: str

@dataclass
class Process:
    node: str
    src: Message
    dst: Message
    start: float
    end: float
    servicetime: float
    id: str
    ue: str

    def __init__(self, src, dst):
        self.src=src
        self.dst=dst
        self.start= src.in_time
        self.end= dst.in_time
        self.servicetime=dst.in_time - src.in_time
        self.node=src.dst
        self.id=src.id
        self.ue = src.ue

class Parser:
    
    def __init__(self, path: str, outfile: str, run: int):
        self.path=path
        self.dataframe = pd.read_csv(path, sep=";")
        print(path)
        self.dataframe = self.dataframe.sort_values(['time'])
        self.messages = []
        self.processes = []
        self.outfile=outfile
        self.run=run
        self._parse_messages()

    def __eq__(self, other):
        """Overrides the default implementation"""
        if isinstance(other, Message):
            if(self.src == other.src and self.dst == other.dst and self.in_time == other.in_time and self.out_time == other.out_time and self.info == other.info):
                return True
        return False
    
    def _check_state_machine(self):
        filtered_messages=[]
        responses=Responses()
        last_message=None
        expected_rsp=None
        next_step=None
        for message in self.messages:
            #print("Message: "+str(message.id))
            if last_message is None and message.id=="InitialUeMessage":
                last_message=message.id
                expected_rsp=responses.responses[last_message]
                next_step=responses.auth_steps[last_message]
                #print("Next: "+str(next_step))
                #print("Expected: "+str(expected_rsp))
                filtered_messages.append(message)
                continue
            elif message.id == next_step:
                last_message=message.id
                if 'GET' in last_message or 'PUT' in last_message or 'PUSH' in last_message or 'POST' in last_message or 'DELETE' in last_message:
                    last_message_http='HTTP'
                else:
                    last_message_http=last_message
                #print("Error:" +last_message_http)
                expected_rsp=responses.responses[last_message_http]
                next_step=responses.auth_steps[last_message]
                #print("Next: "+str(next_step))
                #print("Expected: "+str(expected_rsp))
                filtered_messages.append(message)
                continue
            elif message.id in list(responses.auth_steps.keys()):
                last_message=message.id
                next_step=responses.auth_steps[last_message]
                #print("Next: "+str(next_step))
                #print("Expected: "+str(expected_rsp))
                filtered_messages.append(message)
                if next_step=='Done':
                    break
                continue
            elif message.id in list(responses.responses.keys()) and message.id not in list(responses.auth_steps.keys()):
                last_message=message.id
                if 'GET' in last_message or 'PUT' in last_message or 'PUSH' in last_message or 'POST' in last_message or 'DELETE' in last_message:
                    last_message_http='HTTP'
                else:
                    last_message_http=last_message
                expected_rsp=responses.responses[last_message_http]
                #print("Next: "+str(next_step))
                #print("Expected: "+str(expected_rsp))
                filtered_messages.append(message)
                if expected_rsp=='Done':
                        break
                continue
            elif expected_rsp is not None and next_step is not None and message.id in expected_rsp and next_step!="Done" and expected_rsp!="Done":
                filtered_messages.append(message)
                continue
            elif '200' in message.id or '201' in message.id or '204' in message.id:
                filtered_messages.append(message)
                continue
            else:
                #print("Message not part of the state machine")
                #print(message)
                continue
        self.messages=filtered_messages
            

    def _parse_messages(self):
        rows = self.dataframe.to_dict("records")
        for i in range(0, len(rows)):
            row = rows[i]
            src = row["src"]
            dst = row["dst"]
            in_time = row["time"]
            out_time = row["time"]
            info = row["message"]
            id=""
            ue = row["ue"]
            self.messages.append(Message(src,dst,in_time,out_time,info,id, ue))
        self._aggregate_messages()
        self._calculate_ids()
        self._check_state_machine()
        self._calculate_servicetimes()

    
    def _aggregate_messages(self):
        index_to_remove=[]
        for i in range(0, len(self.messages)-2):
            message = self.messages[i]
            other = self.messages[i+1]
            #print('M SRC {} | O SRC  "{}" _ M DST  "{}" | O DST  "{}"'.format(message.src, other.src,message.dst, other.dst))
            if message.src == other.src and message.dst == other.dst and other.info=="DATA Json data field":
                message.in_time=message.in_time
                message.out_time=other.out_time
                index_to_remove.append(i+1)
                i+=1
        self._remove_messages(index_to_remove)
        index_to_remove=[]
        for i in range(0, len(self.messages)-2):
            message = self.messages[i]
            if "201" in message.info or "200" in message.info:
                for j in range(i+1, len(self.messages)-1):
                    other = self.messages[j]
                    if message.src == other.src and message.dst == other.dst:
                        if other.info=="DATA Json data field":
                            message.in_time=message.in_time
                            message.out_time=other.out_time
                            index_to_remove.append(j)
                        break
        self._remove_messages(index_to_remove)

    def _calculate_servicetimes(self):
        endtime=0
        endue=""
        end_ind=0
        for i in range(0, len(self.messages)-1):
                message = self.messages[i]
                if message.id =="PDUSessionResourceSetup_Rsp":
                    end_ind=i
            #TODO break after last PDU Req in file
        if end_ind >0:
            self.messages=self.messages[:end_ind+2]
        for i in range(0, len(self.messages)-1):
            message = self.messages[i]
            messages_in_process = []
            indices_to_remove = []
            messages_in_process.append(message)
            indices_to_remove.append(i)
            for j in range(i+1, len(self.messages)):
                if(self.messages[j].src == message.dst):
                    messages_in_process.append(self.messages[j])
                    indices_to_remove.append(j)
                    break
            self.processes.append(Process(messages_in_process[0], messages_in_process[-1]))
            endtime=messages_in_process[-1].out_time
            endue=messages_in_process[-1].ue
        if(len(self.processes)>0):
            end_msg=Message(src="none", dst="none", in_time=endtime,out_time=endtime, info="end", id="end",ue=endue)
            self.processes.append(Process(end_msg,end_msg))
        self._export_processes(self.path)
    
    def _remove_messages(self,messages_to_remove):
        for index in sorted(messages_to_remove, reverse=True):
            del self.messages[index]
            

    def _export_processes(self, path):
        csv_process=""
        for process in self.processes:
            csv_process+="\n"+process.node+";"+str(process.servicetime)+";"+str(process.id)+";"+str(self.run)+";"+str(process.start)+";"+str(process.end)+ ";"+ str(process.ue)
        with open(self.outfile, "a") as f:
            f.write(csv_process)
    
    def _calculate_ids(self):
        for message in self.messages:
            info_n=""
            if (message.src == "gnB") | (message.dst == "gnB"):
                for part in message.info.split(","):
                    if ("SACK" not in part) and ("Arwnd" not in part):
                        info_n= info_n+part
                info_n=info_n.split(".")[-1]
                if "PDUSessionResourceSetup" in info_n:
                    if message.src=="gnB":
                        info_n=info_n+"_Rsp"
                    if message.dst=="gnB":
                        info_n=info_n+"_Req"
            elif "DATA" in message.info:
                info_n=message.info 
            elif "PFCP" in message.info:
                info_n=message.info.split(".")[-1]
            else:
                if len(message.info.split(":"))<=1:
                    info_n=message.info.split("[")[0]
                else:
                    data_string=message.info.split(":")[1]
                    if (data_string == " 200") | (data_string == " 201") | (data_string == " 204"):
                        info_n=data_string
                    else: 
                        last_part=data_string.split("/")[-1].split("?")[0]
                        if last_part.isdigit():
                            last_part=""
                        info_n=message.info.split(":")[0]+"-"+data_string.split("/")[0]+data_string.split("/")[1]+data_string.split("/")[2]+last_part
            message.id=str(info_n).replace(" ", "")
def main():
    argParser = argparse.ArgumentParser()
    argParser.add_argument("-p", "--path", help="specify path for files",default='')
    argParser.add_argument("-o", "--outfile", help="specify nf to restart",default='test_multiple.csv')
    args = argParser.parse_args()
    path = args.path
    outfile = args.outfile
    header="node;servicetime;id;run;start;end;ue"
    #with open(outfile, "w") as f:
        #f.write(header)
    for filename in os.listdir(path):
        f = os.path.join(path, filename)
        if os.path.isfile(f) and fnmatch.fnmatch(filename, '*filtered_imsi*_m_stats.csv'):
            with open(f,"r") as file:
                if '999040000000000' in file.read():
                    cleanup_compromised(path, filename)
                    continue #TODO remove everything until next InitialUEMsg 
            run=int(filename.split("_")[6])
            responses=Responses()
            #try:
            parser = Parser(f,outfile,run)
            #except Exception as e:
            #    print(e)
            #    print(f)


def cleanup_compromised(path: str, filename: str):
    TAG = 'InitialUeMessage'
    tag_found = 0
    f = os.path.join(path, filename)
    with open(f,"r") as file:
        lines = file.readlines()
    with open(f,"w") as file:
        file.write("time;src;dst;message;ue\n")
    with open(f, "a") as file:
        for line in lines:
            if TAG in line:
                tag_found+=1
            if tag_found==2:
                file.write(line)

if __name__ == "__main__":
    main()