import pyshark
from collections import defaultdict

FIELDS = (
    "srcaddr",
    "dstaddr",
    "octets",
)


def parse_cflow_packet(packet_cflow):
    """
    Given the cflow layer of a packet (packet.cflow), parse out the desired fields
    and combine them into a single dictionary.

    Output will be a list of dicts like so:

    [
        {"srcaddr": "192.168.1.10", "dstaddr": "192.168.2.10", "octets": "2562"},
        {"srcaddr": "192.168.1.10", "dstaddr": "192.168.2.10", "octets": "270"},
    ]
    """
    flows = []
    for field_idx, field in enumerate(FIELDS):
        field_exists = hasattr(packet_cflow, field)
        if not field_exists:
            continue

        # Use the first field to define the flows
        if field_idx == 0:
            for flow_number, value in enumerate(
                getattr(packet_cflow, field).all_fields
            ):
                flows.append({field: value.showname_value})
            continue

        # Add the additional metadata to their respective flows
        for flow_number, value in enumerate(getattr(packet_cflow, field).all_fields):
            flows[flow_number][field] = value.showname_value

    return flows


capture = pyshark.FileCapture("netflows.pcap")

all_flows = []

for packet in capture:
    packet_flows = parse_cflow_packet(packet.cflow)
    all_flows.extend(packet_flows)


# Calculate total bytes for each unique src/dest pair
flow_octets = defaultdict(int)
for flow in all_flows:
    uid = flow["srcaddr"] + "-" + flow["dstaddr"]
    octets = int(flow["octets"])
    flow_octets[uid] += octets

print(flow_octets)