pkt_reconstruct.py

import os
import argparse
from malware import utils

parser = argparse.ArgumentParser(description='''This is a packet reconstruct
                                 tool to help reconstruct
                                 the packet payload.''')
parser.add_argument("-d", "--directory", type=str,
                    help="Specify a path which place pcap file")
args = parser.parse_args()


def get_pcap_list(path):
    pcap_list = []
    dirs = os.listdir(path)
    dirs.sort()
    for item in dirs:
        # if item.split('.')[-1] == 'pcap':
        pcap_list.append(item)
    return pcap_list

if __name__ == '__main__':
    pcap_list = get_pcap_list(args.directory)
    for pcap in pcap_list:
        save_path = './{log}/{path}/'.format(log=args.directory,
                                             path=pcap)
        pcap_path = './{log}/{path}/{path2}.pcap'.format(log=args.directory,
                                                         path=pcap, path2=pcap)
        connection = utils.follow_tcp_stream(pcap_path)
        utils.dump_tcp_stream_content(connection, save_path, True)

        udp_connection = utils.follow_udp_stream(pcap_path)
        utils.dump_udp_stream_content(udp_connection, save_path, True)