1#!/usr/bin/python 2# Copyright 2012 Google, Inc. All rights reserved. 3 4"""TestCreator creates test templates from pcap files.""" 5 6import argparse 7import base64 8import glob 9import re 10import string 11import subprocess 12import sys 13 14 15class Packet(object): 16 """Helper class encapsulating packet from a pcap file.""" 17 18 def __init__(self, packet_lines): 19 self.packet_lines = packet_lines 20 self.data = self._DecodeText(packet_lines) 21 22 @classmethod 23 def _DecodeText(cls, packet_lines): 24 packet_bytes = [] 25 # First line is timestamp and stuff, skip it. 26 # Format: 0x0010: 0000 0020 3aff 3ffe 0000 0000 0000 0000 ....:.?......... 27 28 for line in packet_lines[1:]: 29 m = re.match(r'\s+0x[a-f\d]+:\s+((?:[\da-f]{2,4}\s)*)', line, re.IGNORECASE) 30 if m is None: continue 31 for hexpart in m.group(1).split(): 32 packet_bytes.append(base64.b16decode(hexpart.upper())) 33 return ''.join(packet_bytes) 34 35 def Test(self, name, link_type): 36 """Yields a test using this packet, as a set of lines.""" 37 yield '// testPacket%s is the packet:' % name 38 for line in self.packet_lines: 39 yield '// ' + line 40 yield 'var testPacket%s = []byte{' % name 41 data = list(self.data) 42 while data: 43 linebytes, data = data[:16], data[16:] 44 yield ''.join(['\t'] + ['0x%02x, ' % ord(c) for c in linebytes]) 45 yield '}' 46 yield 'func TestPacket%s(t *testing.T) {' % name 47 yield '\tp := gopacket.NewPacket(testPacket%s, LinkType%s, gopacket.Default)' % (name, link_type) 48 yield '\tif p.ErrorLayer() != nil {' 49 yield '\t\tt.Error("Failed to decode packet:", p.ErrorLayer().Error())' 50 yield '\t}' 51 yield '\tcheckLayers(p, []gopacket.LayerType{LayerType%s, FILL_ME_IN_WITH_ACTUAL_LAYERS}, t)' % link_type 52 yield '}' 53 yield 'func BenchmarkDecodePacket%s(b *testing.B) {' % name 54 yield '\tfor i := 0; i < b.N; i++ {' 55 yield '\t\tgopacket.NewPacket(testPacket%s, LinkType%s, gopacket.NoCopy)' % (name, link_type) 56 yield '\t}' 57 yield '}' 58 59 60 61def GetTcpdumpOutput(filename): 62 """Runs tcpdump on the given file, returning output as string.""" 63 return subprocess.check_output( 64 ['tcpdump', '-XX', '-s', '0', '-n', '-r', filename]) 65 66 67def TcpdumpOutputToPackets(output): 68 """Reads a pcap file with TCPDump, yielding Packet objects.""" 69 pdata = [] 70 for line in output.splitlines(): 71 if line[0] not in string.whitespace and pdata: 72 yield Packet(pdata) 73 pdata = [] 74 pdata.append(line) 75 if pdata: 76 yield Packet(pdata) 77 78 79def main(): 80 class CustomHelpFormatter(argparse.ArgumentDefaultsHelpFormatter): 81 def _format_usage(self, usage, actions, groups, prefix=None): 82 header =('TestCreator creates gopacket tests using a pcap file.\n\n' 83 'Tests are written to standard out... they can then be \n' 84 'copied into the file of your choice and modified as \n' 85 'you see.\n\n') 86 return header + argparse.ArgumentDefaultsHelpFormatter._format_usage( 87 self, usage, actions, groups, prefix) 88 89 parser = argparse.ArgumentParser(formatter_class=CustomHelpFormatter) 90 parser.add_argument('--link_type', default='Ethernet', help='the link type (default: %(default)s)') 91 parser.add_argument('--name', default='Packet%d', help='the layer type, must have "%d" inside it') 92 parser.add_argument('files', metavar='file.pcap', type=str, nargs='+', help='the files to process') 93 94 args = parser.parse_args() 95 96 for arg in args.files: 97 for path in glob.glob(arg): 98 for i, packet in enumerate(TcpdumpOutputToPackets(GetTcpdumpOutput(path))): 99 print '\n'.join(packet.Test( 100 args.name % i, args.link_type)) 101 102if __name__ == '__main__': 103 main() 104