#!/usr/bin/env python

import config
import dpkt
import sys

def parse_pcap(filename):
    filestream = open(filename, 'rb')
    first_timestamp = None
    dst_port = None
    pcap_stream = dpkt.pcap.Reader(filestream)
    decode = {8: decode_octet_stream, 9: decode_nyte_stream}[config.BYTE_WIDTH]
    for timestamp, packet in pcap_stream:
        stream = dpkt.ethernet.Ethernet(packet).ip.data
        if first_timestamp is None:
            first_timestamp = timestamp
        if dst_port is None:
            dst_port = stream.dport
        if stream.data == '':
            continue
        direction = 'send' if stream.dport == dst_port else 'recv'
        yield {'direction': direction, 'timediff': timestamp-first_timestamp, 'data': decode(stream.data)}

def decode_octet_stream(data):
    return map(ord, data)

def decode_nyte_stream(n):
    bin_str = nytes_to_bit_string(n)
    return [int(bin_str[i:i+9], 2) for i in xrange(0, len(bin_str), 9)]

def nytes_to_bit_string(n):
    bin_str = "".join(bin(ord(c))[2:].zfill(8) for c in n)
    num_bits = (len(n) * 8) % 9
    return bin_str[:len(bin_str) - num_bits]

def get_streams(filename):
    sent = []
    recv = []
    for thingy in parse_pcap(filename):
        if thingy['direction'] == 'send':
            sent.extend(thingy['data'])
        else:
            recv.extend(thingy['data'])
    return sent, recv

def print_streams(filename):
    sent = []
    recv = []
    
    print '''#!/usr/bin/env python

from pwn import *
from pwntools_functionality import *

pp = []

'''    
    for thingy in parse_pcap(filename):
        if thingy['direction'] == 'send':
            print "# send"
            print "pp.append("+repr(thingy['data'])+")"
        else:
            print "# recv"
            print "# " + repr(thingy['data'])
        print ""

def main(filename):
    sent, recv = get_streams(filename)
    if not all(c < 0x7f for c in sent + recv):
        print 'WARNING: unprintable characters have been replaced with "?"'
        print

    print '=== Sent data ==='
    print ''.join(chr(c) if c < 0x7f else '?' for c in sent)
    print
    print '=== Received data ==='
    print ''.join(chr(c) if c < 0x7f else '?' for c in recv)
    print

if __name__ == '__main__':
    if len(sys.argv) == 3:
        if sys.argv[1] == "print":
            main(sys.argv[2])
        elif sys.argv[1] == "print_exploit":
            print_streams(sys.argv[2])
    else:
        print 'Usage: pcap.py print|print_exploit filename.pcap'