1
\$\begingroup\$

I've created a short Python script in order to do some analysis on network flows.

I have an inventory of networks (/8 to /30 masks) of medium size (around 10k references) that provide several information for each of my networks (physical sites, VRF (routing instances), etc.).

I then have some humongous flows data (millions of flows) that contains precise (/32 masks) source and destination IP.

My main goal is to identify which network my /32 IP belongs to, and then to identify the flows where my destination VRF is the same as my source VRF.

import xlsxwriter
import csv
from progress.bar import Bar
from ipaddress import IPv4Address, IPv4Network
import operator
import timeit
import os

AnalysedDictionnary = dict()

class Flow:
    def __init__(self, srcIp : IPv4Address, dstIp : IPv4Address, port : int, nbHits : int):
        self.srcIp = srcIp
        self.dstIp = dstIp
        self.port = port
        self.nbHits = nbHits

    def __repr__(self):
        return "<Flow srcIp:%s dstIp:%s port:%s nbHits:%s>" % (self.srcIp, self.dstIp, self.port, self.nbHits)

    def __str__(self):
        return "From str method of Flow: srcIp is %s, dstIp is %s, port is %s, nbHits is %s," % (self.srcIp, self.dstIp, self.port, self.nbHits)

class EntityNetwork:
    def __init__(self, siteId, siteName, vrfId, vrfName, IPV4Networks):
        self.siteId = siteId
        self.siteName = siteName
        self.vrfId = vrfId
        self.vrfName = vrfName
        self.MyIPv4Network = IPV4Networks

class InventoryNetworks:
    # "Reference all current Known Networks"
    def __init__(self, filePath=''):
        self.filePath = filePath
        self.MyEntityNetwork = []
        # "Read File, parse file and populate a list of IPV4Networks"
        with open(filePath) as csvfile:
            reader = csv.DictReader(csvfile, delimiter=',')
            for line in reader:
                self.MyEntityNetwork.append(EntityNetwork(line['Site ID'], line['Site'], line['VRF (VRF ID)'], line['VRF Description'], IPv4Network(line['Subnet'])))
        self.MyEntityNetwork.sort(key=operator.attrgetter('MyIPv4Network'), reverse=True)
            

class FlowCapture:
    # An observed flow
    def __init__(self, filePath=''):
        self.filePath = filePath
        self.myFlows = []
        # "Read File, parse file and populate a list of IPV4Networks"
        file = open(filePath, "r")
        Lines = file.readlines() 
        for line in Lines: 
            splitedLine = line.split(",")
            self.myFlows.append(Flow(IPv4Address(splitedLine[7]), IPv4Address(splitedLine[8]), splitedLine[25], 1))
        tmpMyFlows = set(self.myFlows)
        self.myFlows = list(tmpMyFlows)

class AnalyzedFlow:
    def findSubnet(self, subnet):
        maxMask = 0

        #Check if we already analyzed the network and saved it in the dictionnary
        if subnet in AnalysedDictionnary:
            return AnalysedDictionnary[subnet]

        for flow in self.refNetworks.MyEntityNetwork:
            if flow.MyIPv4Network.overlaps(IPv4Network(IPv4Address(subnet))):
                if flow.MyIPv4Network.prefixlen > maxMask:
                    matchedSubnet = flow
                    maxMask = flow.MyIPv4Network.prefixlen
                    AnalysedDictionnary[subnet] = matchedSubnet
                    #we ordered the subnets by mask size, meaning that we can't find a more precise mask that would match our network, so we can break
                    break
        return matchedSubnet
                    

    def __init__(self, flow: Flow, refNetworks : InventoryNetworks):
        self.srcIp = flow.srcIp
        self.dstIp = flow.dstIp
        self.port = flow.port
        self.nbHits = flow.nbHits
        self.refNetworks = refNetworks
        self.srcSubnet = self.findSubnet(self.srcIp)
        self.dstSubnet = self.findSubnet(self.dstIp)


#list of interresting ports
adPort = ['10','12','14','443']

#Create our Inventory of Networks from reference file
script_dir = os.path.dirname(os.path.abspath(__file__))
rel_path = "InventorySheet.csv"
abs_file_path = os.path.join(script_dir, rel_path)
ReferenceNetworks = InventoryNetworks(abs_file_path)

#Create our flows
rel_path = "Logs_Equipment.txt"
abs_file_path = os.path.join(script_dir, rel_path)
CapturedFlows = FlowCapture(abs_file_path)

#Create our analyzed flows
AnalysedFlows = []
print('its going on')
start_time = timeit.default_timer()

# code you want to evaluate
#[:100] limit to 100 objects in our array so testing time is shorter
bar = Bar('Processing', max=len(CapturedFlows.myFlows[:100]))
for flow in CapturedFlows.myFlows[:100]:
    AnalysedFlows.append(AnalyzedFlow(flow,ReferenceNetworks))
    bar.next()
bar.finish()
elapsed = timeit.default_timer() - start_time
print('we analysed in :', elapsed)

#Filter our analyzed flow
intraVrfFlows = [flow for flow in AnalysedFlows if flow.srcSubnet.vrfId == flow.dstSubnet.vrfId]

adFlows = [flow for flow in AnalysedFlows if flow.port in adPort]

# Create a workbook and add a worksheet.
workbook = xlsxwriter.Workbook('flowAnalysis.xlsx')
worksheetIntraVrfFlows = workbook.add_worksheet('IntraVrfFlows')
worksheetAdFlows = workbook.add_worksheet('AdFlows')

# Start from the first cell. Rows and columns are zero indexed.
def writeSheetIndex(excelSheet):
    excelSheet.write(0,0, 'Source IP')
    excelSheet.write(0,1, 'Source Subnet')
    excelSheet.write(0,2, 'Source VRF')
    excelSheet.write(0,3, 'Source Site')
    excelSheet.write(0,4, 'Port')
    excelSheet.write(0,5, 'Destination IP')
    excelSheet.write(0,6, 'Destination Subnet')
    excelSheet.write(0,7, 'Destination VRF')
    excelSheet.write(0,8, 'Destination Site')

def populateSheetFlow(excelSheet, analyzedFlows):
    nrow = 1
    for flow in analyzedFlows:
        excelSheet.write(nrow,0, str(flow.srcIp))
        excelSheet.write(nrow,1, flow.srcSubnet.MyIPv4Network.with_prefixlen)
        excelSheet.write(nrow,2, flow.srcSubnet.vrfName)
        excelSheet.write(nrow,3, flow.srcSubnet.siteName)
        excelSheet.write(nrow,4, flow.port)
        excelSheet.write(nrow,5, str(flow.dstIp))
        excelSheet.write(nrow,6, flow.dstSubnet.MyIPv4Network.with_prefixlen)
        excelSheet.write(nrow,7, flow.dstSubnet.vrfName)
        excelSheet.write(nrow,8, flow.dstSubnet.siteName)
        if flow.srcSubnet.siteName :
            nrow = nrow + 1


if True :
    writeSheetIndex(worksheetIntraVrfFlows)
    populateSheetFlow(worksheetIntraVrfFlows, intraVrfFlows)
    writeSheetIndex(worksheetAdFlows)
    populateSheetFlow(worksheetAdFlows, adFlows)
    workbook.close()

The first 2 great improvement I've found are:

  • Order my networks by mask size so that I can stop searching as soon as I have a match
  • Create a dictionary where I save all my IP//Network match so that I don't have to iterate on my entire network inventory if I've already done the job once for this IP

Are there any other places where I could improve my script's performance? So far my metrics are about 6 for 100 rows, 30s for 1k rows, 120s for 10k rows (the more rows at once, the more my dictionary is useful).

That's not too terrible, but I dread how long it will take me to analyze 100 millions of rows, so I'd rather find all optimisations I can beforehand.

\$\endgroup\$
3
\$\begingroup\$

I think the next thing to investigate is this section, inside your AnalyzedFlow class findSubnet method:

for flow in self.refNetworks.MyEntityNetwork:
        if flow.MyIPv4Network.overlaps(IPv4Network(IPv4Address(subnet))):
            if flow.MyIPv4Network.prefixlen > maxMask:

First, you're constructing the IPv4Network(IPv4Address(subnet)) once per loop iteration. Instead, do something like

subnetNetwork = IPv4Network(IPv4Address(subnet))
for flow in self.refNetworks.MyEntityNetwork:
    if flow.MyIPv4Network.overlaps(subnetNetwork):
        if flow.MyIPv4Network.prefixlen > maxMask:

The second bit is more complex. Your flows are in order, but if you structured them differently it might - depending upon your subnets - speed things up.

For example, let's say your 10,000 networks have a few hundred networks all matching the pattern 10.x.y.z/some-mask, a few hundred networks all matching the pattern 11.x.y.z/some-mask, a few hundred networks all matching the pattern 183.x.y.z/some-mask, and so forth. If your input subnet is 183.22.15.4 ideally you don't want to check any of the 10... or 11... networks, just the 'bucket' of networks that start with 183.

It might take a while to structure it, and I can try to help if you need it, but I would build a dict of subnets based on the first octet in the network addresses, so your InventoryNetworks instead of having a list would have a dict.

class InventoryNetworks:
    # "Reference all current Known Networks"
    def __init__(self, filePath=''):
        self.filePath = filePath
        self.MyEntityNetwork = dict()
        # "Read File, parse file and populate a dict of IPV4Networks"
        with open(filePath) as csvfile:
            reader = csv.DictReader(csvfile, delimiter=',')
            for line in reader:
               entityNetwork = EntityNetwork(line['Site ID'], line['Site'], line['VRF (VRF ID)'], line['VRF Description'], IPv4Network(line['Subnet']))
               exploded = entityNetwork.MyIPv4Network.network_address.exploded
               firstOctet = exploded[0:exploded.index(".")]
               prev = MyEntityNetwork[firstOctet]
               if prev is None:
                   MyEntityNetwork[firstOctet] = [ entityNetwork ]
               else:
                   prev.append(entityNetwork)
                   MyEntityNetwork[firstOctet] = prev
           for octet in MyEntityNetwork.keys():
               items = MyEntityNetwork[octet]
               items.sort(key=operator.attrgetter('MyIPv4Network'), reverse=True)
               MyEntityNetwork[octet] = items
        

Then further down you can do something similar to only search subnets that match on the first octet, and that should improve the matching speed. I haven't included that code, I'm hopeful you can figure it out if you need it.

You could even nest this further, and have additional levels in your dict, like

MyEntityNetwork['183']['22']['15'] = [ list of items that match ]
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.