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