1## This file is part of Scapy
2## See http://www.secdev.org/projects/scapy for more informations
3## Copyright (C) Philippe Biondi <phil@secdev.org>
4## This program is published under a GPLv2 license
5
6"""
7General utility functions.
8"""
9
10import os,sys,socket,types
11import random,time
12import gzip,zlib
13import re,struct,array,stat
14import subprocess
15
16import warnings
17warnings.filterwarnings("ignore","tempnam",RuntimeWarning, __name__)
18
19from .config import conf
20from .data import MTU
21from .error import log_runtime,log_loading,log_interactive, Scapy_Exception
22from .base_classes import BasePacketList,BasePacket
23
24
25WINDOWS=sys.platform.startswith("win32")
26
27###########
28## Tools ##
29###########
30
31def get_temp_file(keep=False, autoext=""):
32    import tempfile
33    fd, fname  = tempfile.mkstemp(suffix = ".scapy" + autoext)
34    os.close(fd)
35    if not keep:
36        conf.temp_files.append(fname)
37    return fname
38
39def str2bytes(x):
40  """Convert input argument to bytes"""
41  if type(x) is bytes:
42    return x
43  elif type(x) is str:
44    return bytes([ ord(i) for i in x ])
45  else:
46    return str2bytes(str(x))
47
48def chb(x):
49  if type(x) is str:
50    return x
51  else:
52    return chr(x)
53
54def orb(x):
55  if type(x) is str:
56    return ord(x)
57  else:
58    return x
59
60def any2b(x):
61  if type(x) is not str and type(x) is not bytes:
62    try:
63      x=bytes(x)
64    except:
65      x = str(x)
66  if type(x) is str:
67    x = bytes([ ord(i) for i in x ])
68  return x
69
70def sane_color(x):
71    r=""
72    for i in x:
73        j = orb(i)
74        if (j < 32) or (j >= 127):
75            r=r+conf.color_theme.not_printable(".")
76        else:
77            r=r+chb(i)
78    return r
79
80def sane(x):
81    r=""
82    for i in x:
83        if type(x) is str:
84          j = ord(i)
85        else:
86          j = i
87        if (j < 32) or (j >= 127):
88            r=r+"."
89        else:
90            r=r+chb(i)
91    return r
92
93def lhex(x):
94    if type(x) is int:
95        return hex(x)
96    elif type(x) is tuple:
97        return "(%s)" % ", ".join(map(lhex, x))
98    elif type(x) is list:
99        return "[%s]" % ", ".join(map(lhex, x))
100    else:
101        return x
102
103@conf.commands.register
104def hexdump(x):
105    if type(x) is not str and type(x) is not bytes:
106      try:
107        x=bytes(x)
108      except:
109        x = str(x)
110    l = len(x)
111    i = 0
112    while i < l:
113        print("%04x  " % i,end = " ")
114        for j in range(16):
115            if i+j < l:
116                print("%02X" % orb(x[i+j]), end = " ")
117            else:
118                print("  ", end = " ")
119            if j%16 == 7:
120                print("", end = " ")
121        print(" ", end = " ")
122        print(sane_color(x[i:i+16]))
123        i += 16
124
125@conf.commands.register
126def linehexdump(x, onlyasc=0, onlyhex=0):
127    if type(x) is not str and type(x) is not bytes:
128      try:
129        x=bytes(x)
130      except:
131        x = str(x)
132    l = len(x)
133    if not onlyasc:
134        for i in range(l):
135            print("%02X" % orb(x[i]), end = " ")
136        print("", end = " ")
137    if not onlyhex:
138        print(sane_color(x))
139
140def chexdump(x):
141    if type(x) is not str and type(x) is not bytes:
142      try:
143        x=bytes(x)
144      except:
145        x = str(x)
146    print(", ".join(map(lambda x: "%#04x"%orb(x), x)))
147
148def hexstr(x, onlyasc=0, onlyhex=0):
149    s = []
150    if not onlyasc:
151        s.append(" ".join(map(lambda x:"%02x"%orb(x), x)))
152    if not onlyhex:
153        s.append(sane(x))
154    return "  ".join(s)
155
156
157@conf.commands.register
158def hexdiff(x,y):
159    """Show differences between 2 binary strings"""
160    x=any2b(x)[::-1]
161    y=any2b(y)[::-1]
162    SUBST=1
163    INSERT=1
164    d={}
165    d[-1,-1] = 0,(-1,-1)
166    for j in range(len(y)):
167        d[-1,j] = d[-1,j-1][0]+INSERT, (-1,j-1)
168    for i in range(len(x)):
169        d[i,-1] = d[i-1,-1][0]+INSERT, (i-1,-1)
170
171    for j in range(len(y)):
172        for i in range(len(x)):
173            d[i,j] = min( ( d[i-1,j-1][0]+SUBST*(x[i] != y[j]), (i-1,j-1) ),
174                          ( d[i-1,j][0]+INSERT, (i-1,j) ),
175                          ( d[i,j-1][0]+INSERT, (i,j-1) ) )
176
177
178    backtrackx = []
179    backtracky = []
180    i=len(x)-1
181    j=len(y)-1
182    while not (i == j == -1):
183        i2,j2 = d[i,j][1]
184        backtrackx.append(x[i2+1:i+1])
185        backtracky.append(y[j2+1:j+1])
186        i,j = i2,j2
187
188
189
190    x = y = i = 0
191    colorize = { 0: lambda x:x,
192                -1: conf.color_theme.left,
193                 1: conf.color_theme.right }
194
195    dox=1
196    doy=0
197    l = len(backtrackx)
198    while i < l:
199        separate=0
200        linex = backtrackx[i:i+16]
201        liney = backtracky[i:i+16]
202        xx = sum(len(k) for k in linex)
203        yy = sum(len(k) for k in liney)
204        if dox and not xx:
205            dox = 0
206            doy = 1
207        if dox and linex == liney:
208            doy=1
209
210        if dox:
211            xd = y
212            j = 0
213            while not linex[j]:
214                j += 1
215                xd -= 1
216            print(colorize[doy-dox]("%04x" % xd), end = " ")
217            x += xx
218            line=linex
219        else:
220            print("    ", end = " ")
221        if doy:
222            yd = y
223            j = 0
224            while not liney[j]:
225                j += 1
226                yd -= 1
227            print(colorize[doy-dox]("%04x" % yd), end = " ")
228            y += yy
229            line=liney
230        else:
231            print("    ", end = " ")
232
233        print(" ", end = " ")
234
235        cl = ""
236        for j in range(16):
237            if i+j < l:
238                if line[j]:
239                    col = colorize[(linex[j]!=liney[j])*(doy-dox)]
240                    print(col("%02X" % line[j][0]), end = " ")
241                    if linex[j]==liney[j]:
242                        cl += sane_color(line[j])
243                    else:
244                        cl += col(sane(line[j]))
245                else:
246                    print("  ", end = " ")
247                    cl += " "
248            else:
249                print("  ", end = " ")
250            if j == 7:
251                print("", end = " ")
252
253
254        print(" ",cl)
255
256        if doy or not yy:
257            doy=0
258            dox=1
259            i += 16
260        else:
261            if yy:
262                dox=0
263                doy=1
264            else:
265                i += 16
266
267
268crc32 = zlib.crc32
269
270if struct.pack("H",1) == b"\x00\x01": # big endian
271    def checksum(pkt):
272        if len(pkt) % 2 == 1:
273            pkt += b"\0"
274        s = sum(array.array("H", pkt))
275        s = (s >> 16) + (s & 0xffff)
276        s += s >> 16
277        s = ~s
278        return s & 0xffff
279else:
280    def checksum(pkt):
281        if len(pkt) % 2 == 1:
282            pkt += b"\0"
283        s = sum(array.array("H", pkt))
284        s = (s >> 16) + (s & 0xffff)
285        s += s >> 16
286        s = ~s
287        return (((s>>8)&0xff)|s<<8) & 0xffff
288
289def warning(x):
290    log_runtime.warning(x)
291
292def mac2str(mac):
293    #return "".join(map(lambda x: chr(int(x,16)), mac.split(":")))
294    if type(mac) != str:
295      mac = mac.decode('ascii')
296    return b''.join([ bytes([int(i, 16)]) for i in mac.split(":") ])
297
298def str2mac(s):
299    return ("%02x:"*6)[:-1] % tuple(s)
300
301def strxor(x,y):
302    #return "".join(map(lambda i,j:chr(ord(i)^ord(j)),x,y))
303    return bytes([ i[0] ^ i[1] for i in zip(x,y) ] )
304
305# Workarround bug 643005 : https://sourceforge.net/tracker/?func=detail&atid=105470&aid=643005&group_id=5470
306try:
307    socket.inet_aton("255.255.255.255")
308except socket.error:
309    def inet_aton(x):
310        if x == "255.255.255.255":
311            return b"\xff"*4
312        else:
313            return socket.inet_aton(x)
314else:
315    inet_aton = socket.inet_aton
316
317inet_ntoa = socket.inet_ntoa
318try:
319    inet_ntop = socket.inet_ntop
320    inet_pton = socket.inet_pton
321except AttributeError:
322    from scapy.pton_ntop import *
323    log_loading.info("inet_ntop/pton functions not found. Python IPv6 support not present")
324
325
326def atol(x):
327    try:
328        ip = inet_aton(x)
329    except socket.error:
330        ip = inet_aton(socket.gethostbyname(x))
331    return struct.unpack("!I", ip)[0]
332def ltoa(x):
333    return inet_ntoa(struct.pack("!I", x&0xffffffff))
334
335def itom(x):
336    return (0xffffffff00000000>>x)&0xffffffff
337
338def do_graph(graph,prog=None,format='png',target=None,string=False,options=None, figsize = (12, 12), **kargs):
339    """do_graph(graph, prog=conf.prog.dot, format="png",
340         target=None, options=None, string=False):
341    if networkx library is available and graph is instance of Graph, use networkx.draw
342
343    string: if not False, simply return the graph string
344    graph: GraphViz graph description
345    format: output type (svg, ps, gif, jpg, etc.), passed to dot's "-T" option. Ignored if target==None
346    target: filename. If None uses matplotlib to display
347    prog: which graphviz program to use
348    options: options to be passed to prog"""
349
350    from scapy.arch import NETWORKX
351    if NETWORKX:
352        import networkx as nx
353
354    if NETWORKX and isinstance(graph, nx.Graph):
355        nx.draw(graph, with_labels = True, edge_color = '0.75', **kargs)
356    else: # otherwise use dot as in scapy 2.x
357        if string:
358            return graph
359        if prog is None:
360            prog = conf.prog.dot
361
362        if not target or not format:
363            format = 'png'
364        format = "-T %s" % format
365
366        p = subprocess.Popen("%s %s %s" % (prog,options or "", format or ""), shell = True, stdin = subprocess.PIPE, stdout = subprocess.PIPE)
367        w, r = p.stdin, p.stdout
368        w.write(graph.encode('utf-8'))
369        w.close()
370        if target:
371            with open(target, 'wb') as f:
372                f.write(r.read())
373        else:
374            try:
375                import matplotlib.image as mpimg
376                import matplotlib.pyplot as plt
377                plt.figure(figsize = figsize)
378                plt.axis('off')
379                return plt.imshow(mpimg.imread(r, format = format), **kargs)
380
381            except ImportError:
382                warning('matplotlib.image required for interactive graph viewing. Use target option to write to a file')
383
384_TEX_TR = {
385    "{":"{\\tt\\char123}",
386    "}":"{\\tt\\char125}",
387    "\\":"{\\tt\\char92}",
388    "^":"\\^{}",
389    "$":"\\$",
390    "#":"\\#",
391    "~":"\\~",
392    "_":"\\_",
393    "&":"\\&",
394    "%":"\\%",
395    "|":"{\\tt\\char124}",
396    "~":"{\\tt\\char126}",
397    "<":"{\\tt\\char60}",
398    ">":"{\\tt\\char62}",
399    }
400
401def tex_escape(x):
402    s = ""
403    for c in x:
404        s += _TEX_TR.get(c,c)
405    return s
406
407def colgen(*lstcol,**kargs):
408    """Returns a generator that mixes provided quantities forever
409    trans: a function to convert the three arguments into a color. lambda x,y,z:(x,y,z) by default"""
410    if len(lstcol) < 2:
411        lstcol *= 2
412    trans = kargs.get("trans", lambda x,y,z: (x,y,z))
413    while 1:
414        for i in range(len(lstcol)):
415            for j in range(len(lstcol)):
416                for k in range(len(lstcol)):
417                    if i != j or j != k or k != i:
418                        yield trans(lstcol[(i+j)%len(lstcol)],lstcol[(j+k)%len(lstcol)],lstcol[(k+i)%len(lstcol)])
419
420def incremental_label(label="tag%05i", start=0):
421    while True:
422        yield label % start
423        start += 1
424
425#########################
426#### Enum management ####
427#########################
428
429class EnumElement:
430    _value=None
431    def __init__(self, key, value):
432        self._key = key
433        self._value = value
434    def __repr__(self):
435        return "<%s %s[%r]>" % (self.__dict__.get("_name", self.__class__.__name__), self._key, self._value)
436    def __getattr__(self, attr):
437        return getattr(self._value, attr)
438    def __str__(self):
439        return self._key
440    def __eq__(self, other):
441        #return self._value == int(other)
442        return self._value == hash(other)
443    def __hash__(self):
444        return self._value
445
446
447class Enum_metaclass(type):
448    element_class = EnumElement
449    def __new__(cls, name, bases, dct):
450        rdict={}
451        for k,v in dct.items():
452            if type(v) is int:
453                v = cls.element_class(k,v)
454                dct[k] = v
455                rdict[v] = k
456        dct["__rdict__"] = rdict
457        return super(Enum_metaclass, cls).__new__(cls, name, bases, dct)
458    def __getitem__(self, attr):
459        return self.__rdict__[attr]
460    def __contains__(self, val):
461        return val in self.__rdict__
462    def get(self, attr, val=None):
463        return self._rdict__.get(attr, val)
464    def __repr__(self):
465        return "<%s>" % self.__dict__.get("name", self.__name__)
466
467
468
469###################
470## Object saving ##
471###################
472
473
474def export_object(obj):
475    import dill as pickle
476    import base64
477    return base64.b64encode(gzip.zlib.compress(pickle.dumps(obj,4),9)).decode('utf-8')
478
479
480def import_object(obj):
481    import dill as pickle
482    import base64
483#    if obj is None:
484#        obj = sys.stdin.read().strip().encode('utf-8')
485    if obj is str:
486        obj = obj.strip().encode('utf-8')
487    return pickle.loads(gzip.zlib.decompress(base64.b64decode(obj)))
488
489
490def save_object(fname, obj):
491    import dill as pickle
492    pickle.dump(obj,gzip.open(fname,"wb"))
493
494def load_object(fname):
495    import dill as pickle
496    return pickle.load(gzip.open(fname,"rb"))
497
498@conf.commands.register
499def corrupt_bytes(s, p=0.01, n=None):
500    """Corrupt a given percentage or number of bytes from bytes"""
501    s = str2bytes(s)
502    s = array.array("B",s)
503    l = len(s)
504    if n is None:
505        n = max(1,int(l*p))
506    for i in random.sample(range(l), n):
507        s[i] = (s[i]+random.randint(1,255))%256
508    return s.tobytes()
509
510@conf.commands.register
511def corrupt_bits(s, p=0.01, n=None):
512    """Flip a given percentage or number of bits from bytes"""
513    s = str2bytes(s)
514    s = array.array("B",s)
515    l = len(s)*8
516    if n is None:
517        n = max(1, int(l*p))
518    for i in random.sample(range(l), n):
519        s[i//8] ^= 1 << (i%8)
520    return s.tobytes()
521
522
523#############################
524## pcap capture file stuff ##
525#############################
526
527@conf.commands.register
528def wrpcap(filename, pkt, *args, **kargs):
529    """Write a list of packets to a pcap file
530gz: set to 1 to save a gzipped capture
531linktype: force linktype value
532endianness: "<" or ">", force endianness"""
533    with PcapWriter(filename, *args, **kargs) as pcap:
534        pcap.write(pkt)
535
536@conf.commands.register
537def rdpcap(filename, count=-1):
538    """Read a pcap file and return a packet list
539count: read only <count> packets"""
540    with PcapReader(filename) as pcap:
541        return pcap.read_all(count=count)
542
543class RawPcapReader:
544    """A stateful pcap reader. Each packet is returned as bytes"""
545    def __init__(self, filename):
546        self.filename = filename
547        try:
548            if not stat.S_ISREG(os.stat(filename).st_mode):
549              raise IOError("GZIP detection works only for regular files")
550            self.f = gzip.open(filename,"rb")
551            magic = self.f.read(4)
552        except IOError:
553            self.f = open(filename,"rb")
554            magic = self.f.read(4)
555        if magic == b"\xa1\xb2\xc3\xd4": #big endian
556            self.endian = ">"
557            self.reader = _RawPcapOldReader(self.f, self.endian)
558        elif magic == b"\xd4\xc3\xb2\xa1": #little endian
559            self.endian = "<"
560            self.reader = _RawPcapOldReader(self.f, self.endian)
561        elif magic == b"\x0a\x0d\x0d\x0a": #PcapNG
562            self.reader = _RawPcapNGReader(self.f)
563        else:
564            raise Scapy_Exception("Not a pcap capture file (bad magic)")
565
566    def __enter__(self):
567        return self.reader
568
569    def __exit__(self, exc_type, exc_value, tracback):
570        self.close()
571
572    def __iter__(self):
573        return self.reader.__iter__()
574
575    def dispatch(self, callback):
576        """call the specified callback routine for each packet read
577
578        This is just a convienience function for the main loop
579        that allows for easy launching of packet processing in a
580        thread.
581        """
582        for p in self:
583            callback(p)
584
585    def read_all(self,count=-1):
586        """return a list of all packets in the pcap file
587        """
588        res=[]
589        while count != 0:
590            count -= 1
591            p = self.read_packet()
592            if p is None:
593                break
594            res.append(p)
595        return res
596
597    def recv(self, size=MTU):
598        """ Emulate a socket
599        """
600        return self.read_packet(size)[0]
601
602    def fileno(self):
603        return self.f.fileno()
604
605    def close(self):
606        return self.f.close()
607
608    def read_packet(self, size = MTU):
609        return self.reader.read_packet(size)
610
611def align32(n):
612  return n + (4 - n % 4) % 4
613
614class _RawPcapNGReader:
615    def __init__(self, filep):
616        self.filep = filep
617        self.filep.seek(0, 0)
618        self.endian = '<'
619        self.tsresol = 6
620        self.linktype = None
621
622    def __iter__(self):
623        return self
624
625    def __next__(self):
626        """implement the iterator protocol on a set of packets in a pcapng file"""
627        pkt = self.read_packet()
628        if pkt == None:
629            raise StopIteration
630        return pkt
631
632    def read_packet(self, size = MTU):
633        while True:
634            buf = self._read_bytes(4, check = False)
635            if len(buf) == 0:
636                return None
637            elif len(buf) != 4:
638                raise IOError("PacketNGReader: Premature end of file")
639            block_type, = struct.unpack(self.endian + 'i', buf)
640            if block_type == 168627466: #Section Header b'\x0a\x0d\x0d\x0a'
641                self.read_section_header()
642            elif block_type == 1:
643                self.read_interface_description()
644            elif block_type == 6:
645                return self.read_enhanced_packet(size)
646            else:
647                warning("PacketNGReader: Unparsed block type %d/#%x" % (block_type, block_type))
648                self.read_generic_block()
649
650    def _read_bytes(self, n, check = True):
651        buf = self.filep.read(n)
652        if check and len(buf) < n:
653            raise IOError("PacketNGReader: Premature end of file")
654        return buf
655
656    def read_generic_block(self):
657        block_length, = struct.unpack(self.endian + 'I', self._read_bytes(4))
658        self._read_bytes(block_length - 12)
659        self._check_length(block_length)
660
661    def read_section_header(self):
662        buf = self._read_bytes(16)
663        if buf[4:8] == b'\x1a\x2b\x3c\x4d':
664            self.endian = '>'
665        elif buf[4:8] == b'\x4d\x3c\x2b\x1a':
666            self.endian = '<'
667        else:
668            raise Scapy_Exception('Cannot read byte order value')
669        block_length, _, major_version, minor_version, section_length = struct.unpack(self.endian + 'IIHHi', buf)
670        options = self._read_bytes(block_length - 24)
671        if options:
672            opt = self.parse_options(options)
673            for i in opt.keys():
674                if not i & (0b1 << 15):
675                    warning("PcapNGReader: Unparsed option %d/#%x in section header" % (i, i))
676        self._check_length(block_length)
677
678    def read_interface_description(self):
679        buf = self._read_bytes(12)
680        block_length, self.linktype, reserved, self.snaplen = struct.unpack(self.endian + 'IHHI', buf)
681        options = self._read_bytes(block_length - 20)
682        if options:
683            opt = self.parse_options(options)
684            for i in opt.keys():
685                if 9 in opt:
686                    self.tsresol = opt[9][0]
687                elif not i & (0b1 << 15):
688                    warning("PcapNGReader: Unparsed option %d/#%x in enhanced packet block" % (i, i))
689        try:
690            self.LLcls = conf.l2types[self.linktype]
691        except KeyError:
692            warning("RawPcapReader: unknown LL type [%i]/[%#x]. Using Raw packets" % (self.linktype,self.linktype))
693            self.LLcls = conf.raw_layer
694
695        self._check_length(block_length)
696
697    def read_enhanced_packet(self, size = MTU):
698        buf = self._read_bytes(24)
699        block_length, interface, ts_high, ts_low, caplen, wirelen = struct.unpack(self.endian + 'IIIIII', buf)
700        timestamp = (ts_high << 32) + ts_low
701
702        pkt = self._read_bytes(align32(caplen))[:caplen]
703        options = self._read_bytes(block_length - align32(caplen) - 32)
704        if options:
705            opt = self.parse_options(options)
706            for i in opt.keys():
707                if not i & (0b1 << 15):
708                    warning("PcapNGReader: Unparsed option %d/#%x in enhanced packet block" % (i, i))
709        self._check_length(block_length)
710        return pkt[:MTU], (self.parse_sec(timestamp), self.parse_usec(timestamp), wirelen)
711
712    def parse_sec(self, t):
713        if self.tsresol & 0b10000000:
714            return t >> self.tsresol
715        else:
716            return t // pow(10, self.tsresol)
717
718    def parse_usec(self, t):
719        if self.tsresol & 0b10000000:
720            return t & (1 << self.tsresol) - 1
721        else:
722            return t % pow(10, self.tsresol)
723
724    def parse_options(self, opt):
725        buf = opt
726        options = {}
727        while buf:
728            opt_type, opt_len = struct.unpack(self.endian + 'HH', buf[:4])
729            if opt_type == 0:
730                return options
731            options[opt_type] = buf[4:4 + opt_len]
732            buf = buf[ 4 + align32(opt_len):]
733        return options
734
735    def _check_length(self, block_length):
736        check_length, = struct.unpack(self.endian + 'I', self._read_bytes(4))
737        if check_length != block_length:
738            raise Scapy_Exception('Block length values are not equal')
739
740class _RawPcapOldReader:
741    def __init__(self, filep, endianness):
742        self.endian = endianness
743        self.f = filep
744        hdr = self.f.read(20)
745        if len(hdr)<20:
746            raise Scapy_Exception("Invalid pcap file (too short)")
747        vermaj,vermin,tz,sig,snaplen,linktype = struct.unpack(self.endian+"HHIIII",hdr)
748
749        self.linktype = linktype
750        try:
751            self.LLcls = conf.l2types[self.linktype]
752        except KeyError:
753            raise Scapy_Exception("Scapy PcapReader: unknown LL type [%i]/[%#x]" % (self.linktype,self.linktype))
754            #warning("RawPcapReader: unknown LL type [%i]/[%#x]. Using Raw packets" % (self.linktype,self.linktype))
755            self.LLcls = conf.raw_layer
756
757    def __iter__(self):
758        return self
759
760    def __next__(self):
761        """implement the iterator protocol on a set of packets in a pcap file"""
762        pkt = self.read_packet()
763        if pkt == None:
764            raise StopIteration
765        return pkt
766
767    def read_packet(self, size=MTU):
768        """return a single packet read from the file
769        bytes, (sec, #timestamp seconds
770                usec, #timestamp microseconds
771                wirelen) #actual length of packet
772        returns None when no more packets are available
773        """
774        hdr = self.f.read(16)
775        if len(hdr) < 16:
776            return None
777        sec,usec,caplen,wirelen = struct.unpack(self.endian+"IIII", hdr)
778        s = self.f.read(caplen)[:MTU]
779        return s,(sec,usec,wirelen) # caplen = len(s)
780
781
782class PcapReader(RawPcapReader):
783    def __init__(self, filename):
784        RawPcapReader.__init__(self, filename)
785    def __enter__(self):
786        return self
787    def __iter__(self):
788        return self
789    def __next__(self):
790        """implement the iterator protocol on a set of packets in a pcap file"""
791        pkt = self.read_packet()
792        if pkt == None:
793            raise StopIteration
794        return pkt
795    def read_packet(self, size=MTU):
796        rp = RawPcapReader.read_packet(self,size)
797        if rp is None:
798            return None
799        s,(sec,usec,wirelen) = rp
800
801
802        try:
803            p = self.reader.LLcls(s)
804        except KeyboardInterrupt:
805            raise
806        except:
807            if conf.debug_dissector:
808                raise
809            p = conf.raw_layer(s)
810        p.time = sec+0.000001*usec
811        p.wirelen = wirelen
812        return p
813
814    def read_all(self,count=-1):
815        res = RawPcapReader.read_all(self, count)
816        import scapy.plist
817        return scapy.plist.PacketList(res,name = os.path.basename(self.filename))
818    def recv(self, size=MTU):
819        return self.read_packet(size)
820
821
822class RawPcapWriter:
823    """A stream PCAP writer with more control than wrpcap()"""
824    def __init__(self, filename, linktype=None, gz=False, endianness="", append=False, sync=False):
825        """
826        linktype: force linktype to a given value. If None, linktype is taken
827                  from the first writter packet
828        gz: compress the capture on the fly
829        endianness: force an endianness (little:"<", big:">"). Default is native
830        append: append packets to the capture file instead of truncating it
831        sync: do not bufferize writes to the capture file
832        """
833
834        self.linktype = linktype
835        self.header_present = 0
836        self.append=append
837        self.gz = gz
838        self.endian = endianness
839        self.filename=filename
840        self.sync=sync
841        bufsz=4096
842        if sync:
843            bufsz=0
844
845        self.f = [open,gzip.open][gz](filename,append and "ab" or "wb", gz and 9 or bufsz)
846
847    def fileno(self):
848        return self.f.fileno()
849
850    def _write_header(self, pkt):
851        self.header_present=1
852
853        if self.append:
854            # Even if prone to race conditions, this seems to be
855            # safest way to tell whether the header is already present
856            # because we have to handle compressed streams that
857            # are not as flexible as basic files
858            g = [open,gzip.open][self.gz](self.filename,"rb")
859            if g.read(16):
860                return
861
862        self.f.write(struct.pack(self.endian+"IHHIIII", 0xa1b2c3d4,
863                                 2, 4, 0, 0, MTU, self.linktype))
864        self.f.flush()
865
866
867    def write(self, pkt):
868        """accepts a either a single packet or a list of packets
869        to be written to the dumpfile
870        """
871        if not self.header_present:
872            self._write_header(pkt)
873        if type(pkt) is bytes:
874            self._write_packet(pkt)
875        else:
876            for p in pkt:
877                self._write_packet(p)
878
879    def _write_packet(self, packet, sec=None, usec=None, caplen=None, wirelen=None):
880        """writes a single packet to the pcap file
881        """
882        if caplen is None:
883            caplen = len(packet)
884        if wirelen is None:
885            wirelen = caplen
886        if sec is None or usec is None:
887            t=time.time()
888            it = int(t)
889            if sec is None:
890                sec = it
891            if usec is None:
892                usec = int(round((t-it)*1000000))
893        self.f.write(struct.pack(self.endian+"IIII", sec, usec, caplen, wirelen))
894        self.f.write(packet)
895        if self.gz and self.sync:
896            self.f.flush()
897
898    def flush(self):
899        return self.f.flush()
900
901    def close(self):
902        return self.f.close()
903
904    def __enter__(self):
905        return self
906    def __exit__(self, exc_type, exc_value, tracback):
907        self.flush()
908        self.close()
909
910
911class PcapWriter(RawPcapWriter):
912    def _write_header(self, pkt):
913        if self.linktype == None:
914            if type(pkt) is list or type(pkt) is tuple or isinstance(pkt,BasePacketList):
915                pkt = pkt[0]
916            try:
917                self.linktype = conf.l2types[pkt.__class__]
918            except KeyError:
919                warning("PcapWriter: unknown LL type for %s. Using type 1 (Ethernet)" % pkt.__class__.__name__)
920                self.linktype = 1
921        RawPcapWriter._write_header(self, pkt)
922
923    def _write_packet(self, packet):
924        try:
925          sec = int(packet.time)
926          usec = int(round((packet.time-sec)*1000000))
927          s = bytes(packet)
928          caplen = len(s)
929          RawPcapWriter._write_packet(self, s, sec, usec, caplen, caplen)
930        except Exception as e:
931          log_interactive.error(e)
932    def write(self, pkt):
933        """accepts a either a single packet or a list of packets
934        to be written to the dumpfile
935        """
936        if not self.header_present:
937            self._write_header(pkt)
938        if isinstance(pkt, BasePacket):
939          self._write_packet(pkt)
940        else:
941            for p in pkt:
942                self._write_packet(p)
943
944
945re_extract_hexcap = re.compile("^((0x)?[0-9a-fA-F]{2,}[ :\t]{,3}|) *(([0-9a-fA-F]{2} {,2}){,16})")
946
947def import_hexcap():
948    p = ""
949    try:
950        while 1:
951            l = raw_input().strip()
952            try:
953                p += re_extract_hexcap.match(l).groups()[2]
954            except:
955                warning("Parsing error during hexcap")
956                continue
957    except EOFError:
958        pass
959
960    p = p.replace(" ","")
961    return p.decode("hex")
962
963
964
965@conf.commands.register
966def wireshark(pktlist, *args):
967    """Run wireshark on a list of packets"""
968    fname = get_temp_file()
969    wrpcap(fname, pktlist)
970    subprocess.Popen([conf.prog.wireshark, "-r", fname] + list(args))
971
972@conf.commands.register
973def tdecode(pkt, *args):
974    """Run tshark to decode and display the packet. If no args defined uses -V"""
975    if not args:
976      args = [ "-V" ]
977    fname = get_temp_file()
978    wrpcap(fname,[pkt])
979    subprocess.call(["tshark", "-r", fname] + list(args))
980
981@conf.commands.register
982def hexedit(x):
983    """Run external hex editor on a packet or bytes. Set editor in conf.prog.hexedit"""
984    x = bytes(x)
985    fname = get_temp_file()
986    with open(fname,"wb") as f:
987      f.write(x)
988    subprocess.call([conf.prog.hexedit, fname])
989    with open(fname, "rb") as f:
990      x = f.read()
991    return x
992
993def __make_table(yfmtfunc, fmtfunc, endline, items, fxyz, sortx=None, sorty=None, seplinefunc=None):
994    vx = {}
995    vy = {}
996    vz = {}
997    vxf = {}
998    vyf = {}
999    max_length = 0
1000    for record in items:
1001        xx,yy,zz = map(str, fxyz(record[0], record[1]))
1002        max_length = max(len(yy),max_length)
1003        vx[xx] = max(vx.get(xx,0), len(xx), len(zz))
1004        vy[yy] = None
1005        vz[(xx,yy)] = zz
1006
1007    vxk = list(vx.keys())
1008    vyk = list(vy.keys())
1009    if sortx:
1010        vxk.sort(sortx)
1011    else:
1012        try:
1013            vxk.sort(key = lambda x: atol(x))
1014        except:
1015            vxk.sort()
1016    if sorty:
1017        vyk.sort(sorty)
1018    else:
1019        try:
1020            vyk.sort(key = lambda x: atol(x))
1021        except:
1022            vyk.sort()
1023
1024
1025    if seplinefunc:
1026        sepline = seplinefunc(max_length, [vx[x] for x in vxk])
1027        print(sepline)
1028
1029    fmt = yfmtfunc(max_length)
1030    print(fmt % "", end = " ")
1031    for x in vxk:
1032        vxf[x] = fmtfunc(vx[x])
1033        print(vxf[x] % x, end = " ")
1034    print(endline)
1035    if seplinefunc:
1036        print(sepline)
1037    for y in vyk:
1038        print(fmt % y, end = " ")
1039        for x in vxk:
1040            print(vxf[x] % vz.get((x,y), "-"), end = " ")
1041        print(endline)
1042    if seplinefunc:
1043        print(sepline)
1044
1045def make_table(*args, **kargs):
1046    __make_table(lambda l:"%%-%is" % l, lambda l:"%%-%is" % l, "", *args, **kargs)
1047
1048def make_lined_table(*args, **kargs):
1049    __make_table(lambda l:"%%-%is |" % l, lambda l:"%%-%is |" % l, "",
1050                 seplinefunc=lambda max_length,x:"+".join([ "-"*(y+2) for y in [max_length-1]+x+[-2]]),
1051                 *args, **kargs)
1052
1053def make_tex_table(*args, **kargs):
1054    __make_table(lambda l: "%s", lambda l: "& %s", "\\\\", seplinefunc=lambda a,x:"\\hline", *args, **kargs)
1055
1056