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,cPickle
13import re,struct,array
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
23
24WINDOWS=sys.platform.startswith("win32")
25
26###########
27## Tools ##
28###########
29
30def get_temp_file(keep=False, autoext=""):
31    f = os.tempnam("","scapy")
32    if not keep:
33        conf.temp_files.append(f+autoext)
34    return f
35
36def sane_color(x):
37    r=""
38    for i in x:
39        j = ord(i)
40        if (j < 32) or (j >= 127):
41            r=r+conf.color_theme.not_printable(".")
42        else:
43            r=r+i
44    return r
45
46def sane(x):
47    r=""
48    for i in x:
49        j = ord(i)
50        if (j < 32) or (j >= 127):
51            r=r+"."
52        else:
53            r=r+i
54    return r
55
56def lhex(x):
57    if type(x) in (int,long):
58        return hex(x)
59    elif type(x) is tuple:
60        return "(%s)" % ", ".join(map(lhex, x))
61    elif type(x) is list:
62        return "[%s]" % ", ".join(map(lhex, x))
63    else:
64        return x
65
66@conf.commands.register
67def hexdump(x):
68    x=str(x)
69    l = len(x)
70    i = 0
71    while i < l:
72        print "%04x  " % i,
73        for j in range(16):
74            if i+j < l:
75                print "%02X" % ord(x[i+j]),
76            else:
77                print "  ",
78            if j%16 == 7:
79                print "",
80        print " ",
81        print sane_color(x[i:i+16])
82        i += 16
83
84@conf.commands.register
85def linehexdump(x, onlyasc=0, onlyhex=0):
86    x = str(x)
87    l = len(x)
88    if not onlyasc:
89        for i in range(l):
90            print "%02X" % ord(x[i]),
91        print "",
92    if not onlyhex:
93        print sane_color(x)
94
95def chexdump(x):
96    x=str(x)
97    print ", ".join(map(lambda x: "%#04x"%ord(x), x))
98
99def hexstr(x, onlyasc=0, onlyhex=0):
100    s = []
101    if not onlyasc:
102        s.append(" ".join(map(lambda x:"%02x"%ord(x), x)))
103    if not onlyhex:
104        s.append(sane(x))
105    return "  ".join(s)
106
107
108@conf.commands.register
109def hexdiff(x,y):
110    """Show differences between 2 binary strings"""
111    x=str(x)[::-1]
112    y=str(y)[::-1]
113    SUBST=1
114    INSERT=1
115    d={}
116    d[-1,-1] = 0,(-1,-1)
117    for j in range(len(y)):
118        d[-1,j] = d[-1,j-1][0]+INSERT, (-1,j-1)
119    for i in range(len(x)):
120        d[i,-1] = d[i-1,-1][0]+INSERT, (i-1,-1)
121
122    for j in range(len(y)):
123        for i in range(len(x)):
124            d[i,j] = min( ( d[i-1,j-1][0]+SUBST*(x[i] != y[j]), (i-1,j-1) ),
125                          ( d[i-1,j][0]+INSERT, (i-1,j) ),
126                          ( d[i,j-1][0]+INSERT, (i,j-1) ) )
127
128
129    backtrackx = []
130    backtracky = []
131    i=len(x)-1
132    j=len(y)-1
133    while not (i == j == -1):
134        i2,j2 = d[i,j][1]
135        backtrackx.append(x[i2+1:i+1])
136        backtracky.append(y[j2+1:j+1])
137        i,j = i2,j2
138
139
140
141    x = y = i = 0
142    colorize = { 0: lambda x:x,
143                -1: conf.color_theme.left,
144                 1: conf.color_theme.right }
145
146    dox=1
147    doy=0
148    l = len(backtrackx)
149    while i < l:
150        separate=0
151        linex = backtrackx[i:i+16]
152        liney = backtracky[i:i+16]
153        xx = sum(len(k) for k in linex)
154        yy = sum(len(k) for k in liney)
155        if dox and not xx:
156            dox = 0
157            doy = 1
158        if dox and linex == liney:
159            doy=1
160
161        if dox:
162            xd = y
163            j = 0
164            while not linex[j]:
165                j += 1
166                xd -= 1
167            print colorize[doy-dox]("%04x" % xd),
168            x += xx
169            line=linex
170        else:
171            print "    ",
172        if doy:
173            yd = y
174            j = 0
175            while not liney[j]:
176                j += 1
177                yd -= 1
178            print colorize[doy-dox]("%04x" % yd),
179            y += yy
180            line=liney
181        else:
182            print "    ",
183
184        print " ",
185
186        cl = ""
187        for j in range(16):
188            if i+j < l:
189                if line[j]:
190                    col = colorize[(linex[j]!=liney[j])*(doy-dox)]
191                    print col("%02X" % ord(line[j])),
192                    if linex[j]==liney[j]:
193                        cl += sane_color(line[j])
194                    else:
195                        cl += col(sane(line[j]))
196                else:
197                    print "  ",
198                    cl += " "
199            else:
200                print "  ",
201            if j == 7:
202                print "",
203
204
205        print " ",cl
206
207        if doy or not yy:
208            doy=0
209            dox=1
210            i += 16
211        else:
212            if yy:
213                dox=0
214                doy=1
215            else:
216                i += 16
217
218
219crc32 = zlib.crc32
220
221if struct.pack("H",1) == "\x00\x01": # big endian
222    def checksum(pkt):
223        if len(pkt) % 2 == 1:
224            pkt += "\0"
225        s = sum(array.array("H", pkt))
226        s = (s >> 16) + (s & 0xffff)
227        s += s >> 16
228        s = ~s
229        return s & 0xffff
230else:
231    def checksum(pkt):
232        if len(pkt) % 2 == 1:
233            pkt += "\0"
234        s = sum(array.array("H", pkt))
235        s = (s >> 16) + (s & 0xffff)
236        s += s >> 16
237        s = ~s
238        return (((s>>8)&0xff)|s<<8) & 0xffff
239
240def warning(x):
241    log_runtime.warning(x)
242
243def mac2str(mac):
244    return "".join(map(lambda x: chr(int(x,16)), mac.split(":")))
245
246def str2mac(s):
247    return ("%02x:"*6)[:-1] % tuple(map(ord, s))
248
249def strxor(x,y):
250    return "".join(map(lambda x,y:chr(ord(x)^ord(y)),x,y))
251
252# Workarround bug 643005 : https://sourceforge.net/tracker/?func=detail&atid=105470&aid=643005&group_id=5470
253try:
254    socket.inet_aton("255.255.255.255")
255except socket.error:
256    def inet_aton(x):
257        if x == "255.255.255.255":
258            return "\xff"*4
259        else:
260            return socket.inet_aton(x)
261else:
262    inet_aton = socket.inet_aton
263
264inet_ntoa = socket.inet_ntoa
265try:
266    inet_ntop = socket.inet_ntop
267    inet_pton = socket.inet_pton
268except AttributeError:
269    from scapy.pton_ntop import *
270    log_loading.info("inet_ntop/pton functions not found. Python IPv6 support not present")
271
272
273def atol(x):
274    try:
275        ip = inet_aton(x)
276    except socket.error:
277        ip = inet_aton(socket.gethostbyname(x))
278    return struct.unpack("!I", ip)[0]
279def ltoa(x):
280    return inet_ntoa(struct.pack("!I", x&0xffffffff))
281
282def itom(x):
283    return (0xffffffff00000000L>>x)&0xffffffffL
284
285def do_graph(graph,prog=None,format=None,target=None,type=None,string=None,options=None):
286    """do_graph(graph, prog=conf.prog.dot, format="svg",
287         target="| conf.prog.display", options=None, [string=1]):
288    string: if not None, simply return the graph string
289    graph: GraphViz graph description
290    format: output type (svg, ps, gif, jpg, etc.), passed to dot's "-T" option
291    target: filename or redirect. Defaults pipe to Imagemagick's display program
292    prog: which graphviz program to use
293    options: options to be passed to prog"""
294
295    if format is None:
296        if WINDOWS:
297            format = "png" # use common format to make sure a viewer is installed
298        else:
299            format = "svg"
300    if string:
301        return graph
302    if type is not None:
303        format=type
304    if prog is None:
305        prog = conf.prog.dot
306    start_viewer=False
307    if target is None:
308        if WINDOWS:
309            tempfile = os.tempnam("", "scapy") + "." + format
310            target = "> %s" % tempfile
311            start_viewer = True
312        else:
313            target = "| %s" % conf.prog.display
314    if format is not None:
315        format = "-T %s" % format
316    w,r = os.popen2("%s %s %s %s" % (prog,options or "", format or "", target))
317    w.write(graph)
318    w.close()
319    if start_viewer:
320        # Workaround for file not found error: We wait until tempfile is written.
321        waiting_start = time.time()
322        while not os.path.exists(tempfile):
323            time.sleep(0.1)
324            if time.time() - waiting_start > 3:
325                warning("Temporary file '%s' could not be written. Graphic will not be displayed." % tempfile)
326                break
327        else:
328            if conf.prog.display == conf.prog._default:
329                os.startfile(tempfile)
330            else:
331                subprocess.Popen([conf.prog.display, tempfile])
332
333_TEX_TR = {
334    "{":"{\\tt\\char123}",
335    "}":"{\\tt\\char125}",
336    "\\":"{\\tt\\char92}",
337    "^":"\\^{}",
338    "$":"\\$",
339    "#":"\\#",
340    "~":"\\~",
341    "_":"\\_",
342    "&":"\\&",
343    "%":"\\%",
344    "|":"{\\tt\\char124}",
345    "~":"{\\tt\\char126}",
346    "<":"{\\tt\\char60}",
347    ">":"{\\tt\\char62}",
348    }
349
350def tex_escape(x):
351    s = ""
352    for c in x:
353        s += _TEX_TR.get(c,c)
354    return s
355
356def colgen(*lstcol,**kargs):
357    """Returns a generator that mixes provided quantities forever
358    trans: a function to convert the three arguments into a color. lambda x,y,z:(x,y,z) by default"""
359    if len(lstcol) < 2:
360        lstcol *= 2
361    trans = kargs.get("trans", lambda x,y,z: (x,y,z))
362    while 1:
363        for i in range(len(lstcol)):
364            for j in range(len(lstcol)):
365                for k in range(len(lstcol)):
366                    if i != j or j != k or k != i:
367                        yield trans(lstcol[(i+j)%len(lstcol)],lstcol[(j+k)%len(lstcol)],lstcol[(k+i)%len(lstcol)])
368
369def incremental_label(label="tag%05i", start=0):
370    while True:
371        yield label % start
372        start += 1
373
374#########################
375#### Enum management ####
376#########################
377
378class EnumElement:
379    _value=None
380    def __init__(self, key, value):
381        self._key = key
382        self._value = value
383    def __repr__(self):
384        return "<%s %s[%r]>" % (self.__dict__.get("_name", self.__class__.__name__), self._key, self._value)
385    def __getattr__(self, attr):
386        return getattr(self._value, attr)
387    def __str__(self):
388        return self._key
389    def __eq__(self, other):
390        return self._value == int(other)
391
392
393class Enum_metaclass(type):
394    element_class = EnumElement
395    def __new__(cls, name, bases, dct):
396        rdict={}
397        for k,v in dct.iteritems():
398            if type(v) is int:
399                v = cls.element_class(k,v)
400                dct[k] = v
401                rdict[v] = k
402        dct["__rdict__"] = rdict
403        return super(Enum_metaclass, cls).__new__(cls, name, bases, dct)
404    def __getitem__(self, attr):
405        return self.__rdict__[attr]
406    def __contains__(self, val):
407        return val in self.__rdict__
408    def get(self, attr, val=None):
409        return self._rdict__.get(attr, val)
410    def __repr__(self):
411        return "<%s>" % self.__dict__.get("name", self.__name__)
412
413
414
415###################
416## Object saving ##
417###################
418
419
420def export_object(obj):
421    print gzip.zlib.compress(cPickle.dumps(obj,2),9).encode("base64")
422
423def import_object(obj=None):
424    if obj is None:
425        obj = sys.stdin.read()
426    return cPickle.loads(gzip.zlib.decompress(obj.strip().decode("base64")))
427
428
429def save_object(fname, obj):
430    cPickle.dump(obj,gzip.open(fname,"wb"))
431
432def load_object(fname):
433    return cPickle.load(gzip.open(fname,"rb"))
434
435@conf.commands.register
436def corrupt_bytes(s, p=0.01, n=None):
437    """Corrupt a given percentage or number of bytes from a string"""
438    s = array.array("B",str(s))
439    l = len(s)
440    if n is None:
441        n = max(1,int(l*p))
442    for i in random.sample(xrange(l), n):
443        s[i] = (s[i]+random.randint(1,255))%256
444    return s.tostring()
445
446@conf.commands.register
447def corrupt_bits(s, p=0.01, n=None):
448    """Flip a given percentage or number of bits from a string"""
449    s = array.array("B",str(s))
450    l = len(s)*8
451    if n is None:
452        n = max(1,int(l*p))
453    for i in random.sample(xrange(l), n):
454        s[i/8] ^= 1 << (i%8)
455    return s.tostring()
456
457
458
459
460#############################
461## pcap capture file stuff ##
462#############################
463
464@conf.commands.register
465def wrpcap(filename, pkt, *args, **kargs):
466    """Write a list of packets to a pcap file
467gz: set to 1 to save a gzipped capture
468linktype: force linktype value
469endianness: "<" or ">", force endianness"""
470    PcapWriter(filename, *args, **kargs).write(pkt)
471
472@conf.commands.register
473def rdpcap(filename, count=-1):
474    """Read a pcap file and return a packet list
475count: read only <count> packets"""
476    return PcapReader(filename).read_all(count=count)
477
478
479
480class RawPcapReader:
481    """A stateful pcap reader. Each packet is returned as a string"""
482
483    def __init__(self, filename):
484        self.filename = filename
485        try:
486            self.f = gzip.open(filename,"rb")
487            magic = self.f.read(4)
488        except IOError:
489            self.f = open(filename,"rb")
490            magic = self.f.read(4)
491        if magic == "\xa1\xb2\xc3\xd4": #big endian
492            self.endian = ">"
493        elif  magic == "\xd4\xc3\xb2\xa1": #little endian
494            self.endian = "<"
495        else:
496            raise Scapy_Exception("Not a pcap capture file (bad magic)")
497        hdr = self.f.read(20)
498        if len(hdr)<20:
499            raise Scapy_Exception("Invalid pcap file (too short)")
500        vermaj,vermin,tz,sig,snaplen,linktype = struct.unpack(self.endian+"HHIIII",hdr)
501
502        self.linktype = linktype
503
504
505
506    def __iter__(self):
507        return self
508
509    def next(self):
510        """impliment the iterator protocol on a set of packets in a pcap file"""
511        pkt = self.read_packet()
512        if pkt == None:
513            raise StopIteration
514        return pkt
515
516
517    def read_packet(self, size=MTU):
518        """return a single packet read from the file
519
520        returns None when no more packets are available
521        """
522        hdr = self.f.read(16)
523        if len(hdr) < 16:
524            return None
525        sec,usec,caplen,wirelen = struct.unpack(self.endian+"IIII", hdr)
526        s = self.f.read(caplen)[:MTU]
527        return s,(sec,usec,wirelen) # caplen = len(s)
528
529
530    def dispatch(self, callback):
531        """call the specified callback routine for each packet read
532
533        This is just a convienience function for the main loop
534        that allows for easy launching of packet processing in a
535        thread.
536        """
537        for p in self:
538            callback(p)
539
540    def read_all(self,count=-1):
541        """return a list of all packets in the pcap file
542        """
543        res=[]
544        while count != 0:
545            count -= 1
546            p = self.read_packet()
547            if p is None:
548                break
549            res.append(p)
550        return res
551
552    def recv(self, size=MTU):
553        """ Emulate a socket
554        """
555        return self.read_packet(size)[0]
556
557    def fileno(self):
558        return self.f.fileno()
559
560    def close(self):
561        return self.f.close()
562
563    def __enter__(self):
564        return self
565
566    def __exit__(self, exc_type, exc_value, tracback):
567        pass
568
569
570class PcapReader(RawPcapReader):
571    def __init__(self, filename):
572        RawPcapReader.__init__(self, filename)
573        try:
574            self.LLcls = conf.l2types[self.linktype]
575        except KeyError:
576            raise Scapy_Exception("Scapy PcapReader: unknown LL type [%i]/[%#x]" % (self.linktype,self.linktype))
577            #warning("PcapReader: unknown LL type [%i]/[%#x]. Using Raw packets" % (self.linktype,self.linktype))
578            self.LLcls = conf.raw_layer
579    def read_packet(self, size=MTU):
580        rp = RawPcapReader.read_packet(self,size)
581        if rp is None:
582            return None
583        s,(sec,usec,wirelen) = rp
584
585        try:
586            p = self.LLcls(s)
587        except KeyboardInterrupt:
588            raise
589        except:
590            if conf.debug_dissector:
591                raise
592            p = conf.raw_layer(s)
593        p.time = sec+0.000001*usec
594        return p
595    def read_all(self,count=-1):
596        res = RawPcapReader.read_all(self, count)
597        import plist
598        return plist.PacketList(res,name = os.path.basename(self.filename))
599    def recv(self, size=MTU):
600        return self.read_packet(size)
601
602
603
604class RawPcapWriter:
605    """A stream PCAP writer with more control than wrpcap()"""
606    def __init__(self, filename, linktype=None, gz=False, endianness="", append=False, sync=False):
607        """
608        linktype: force linktype to a given value. If None, linktype is taken
609                  from the first writter packet
610        gz: compress the capture on the fly
611        endianness: force an endianness (little:"<", big:">"). Default is native
612        append: append packets to the capture file instead of truncating it
613        sync: do not bufferize writes to the capture file
614        """
615
616        self.linktype = linktype
617        self.header_present = 0
618        self.append=append
619        self.gz = gz
620        self.endian = endianness
621        self.filename=filename
622        self.sync=sync
623        bufsz=4096
624        if sync:
625            bufsz=0
626
627        self.f = [open,gzip.open][gz](filename,append and "ab" or "wb", gz and 9 or bufsz)
628
629    def fileno(self):
630        return self.f.fileno()
631
632    def _write_header(self, pkt):
633        self.header_present=1
634
635        if self.append:
636            # Even if prone to race conditions, this seems to be
637            # safest way to tell whether the header is already present
638            # because we have to handle compressed streams that
639            # are not as flexible as basic files
640            g = [open,gzip.open][self.gz](self.filename,"rb")
641            if g.read(16):
642                return
643
644        self.f.write(struct.pack(self.endian+"IHHIIII", 0xa1b2c3d4L,
645                                 2, 4, 0, 0, MTU, self.linktype))
646        self.f.flush()
647
648
649    def write(self, pkt):
650        """accepts a either a single packet or a list of packets
651        to be written to the dumpfile
652        """
653        if not self.header_present:
654            self._write_header(pkt)
655        if type(pkt) is str:
656            self._write_packet(pkt)
657        else:
658            for p in pkt:
659                self._write_packet(p)
660
661    def _write_packet(self, packet, sec=None, usec=None, caplen=None, wirelen=None):
662        """writes a single packet to the pcap file
663        """
664        if caplen is None:
665            caplen = len(packet)
666        if wirelen is None:
667            wirelen = caplen
668        if sec is None or usec is None:
669            t=time.time()
670            it = int(t)
671            if sec is None:
672                sec = it
673            if usec is None:
674                usec = int(round((t-it)*1000000))
675        self.f.write(struct.pack(self.endian+"IIII", sec, usec, caplen, wirelen))
676        self.f.write(packet)
677        if self.gz and self.sync:
678            self.f.flush()
679
680    def flush(self):
681        return self.f.flush()
682
683    def close(self):
684        return self.f.close()
685
686    def __enter__(self):
687        return self
688    def __exit__(self, exc_type, exc_value, tracback):
689        self.flush()
690
691
692class PcapWriter(RawPcapWriter):
693    def _write_header(self, pkt):
694        if self.linktype == None:
695            if type(pkt) is list or type(pkt) is tuple or isinstance(pkt,BasePacketList):
696                pkt = pkt[0]
697            try:
698                self.linktype = conf.l2types[pkt.__class__]
699            except KeyError:
700                warning("PcapWriter: unknown LL type for %s. Using type 1 (Ethernet)" % pkt.__class__.__name__)
701                self.linktype = 1
702        RawPcapWriter._write_header(self, pkt)
703
704    def _write_packet(self, packet):
705        sec = int(packet.time)
706        usec = int(round((packet.time-sec)*1000000))
707        s = str(packet)
708        caplen = len(s)
709        RawPcapWriter._write_packet(self, s, sec, usec, caplen, caplen)
710
711
712re_extract_hexcap = re.compile("^((0x)?[0-9a-fA-F]{2,}[ :\t]{,3}|) *(([0-9a-fA-F]{2} {,2}){,16})")
713
714def import_hexcap():
715    p = ""
716    try:
717        while 1:
718            l = raw_input().strip()
719            try:
720                p += re_extract_hexcap.match(l).groups()[2]
721            except:
722                warning("Parsing error during hexcap")
723                continue
724    except EOFError:
725        pass
726
727    p = p.replace(" ","")
728    return p.decode("hex")
729
730
731
732@conf.commands.register
733def wireshark(pktlist):
734    """Run wireshark on a list of packets"""
735    f = get_temp_file()
736    wrpcap(f, pktlist)
737    subprocess.Popen([conf.prog.wireshark, "-r", f])
738
739@conf.commands.register
740def hexedit(x):
741    x = str(x)
742    f = get_temp_file()
743    open(f,"w").write(x)
744    subprocess.call([conf.prog.hexedit, f])
745    x = open(f).read()
746    os.unlink(f)
747    return x
748
749def __make_table(yfmtfunc, fmtfunc, endline, list, fxyz, sortx=None, sorty=None, seplinefunc=None):
750    vx = {}
751    vy = {}
752    vz = {}
753    vxf = {}
754    vyf = {}
755    l = 0
756    for e in list:
757        xx,yy,zz = map(str, fxyz(e))
758        l = max(len(yy),l)
759        vx[xx] = max(vx.get(xx,0), len(xx), len(zz))
760        vy[yy] = None
761        vz[(xx,yy)] = zz
762
763    vxk = vx.keys()
764    vyk = vy.keys()
765    if sortx:
766        vxk.sort(sortx)
767    else:
768        try:
769            vxk.sort(lambda x,y:int(x)-int(y))
770        except:
771            try:
772                vxk.sort(lambda x,y: cmp(atol(x),atol(y)))
773            except:
774                vxk.sort()
775    if sorty:
776        vyk.sort(sorty)
777    else:
778        try:
779            vyk.sort(lambda x,y:int(x)-int(y))
780        except:
781            try:
782                vyk.sort(lambda x,y: cmp(atol(x),atol(y)))
783            except:
784                vyk.sort()
785
786
787    if seplinefunc:
788        sepline = seplinefunc(l, map(lambda x:vx[x],vxk))
789        print sepline
790
791    fmt = yfmtfunc(l)
792    print fmt % "",
793    for x in vxk:
794        vxf[x] = fmtfunc(vx[x])
795        print vxf[x] % x,
796    print endline
797    if seplinefunc:
798        print sepline
799    for y in vyk:
800        print fmt % y,
801        for x in vxk:
802            print vxf[x] % vz.get((x,y), "-"),
803        print endline
804    if seplinefunc:
805        print sepline
806
807def make_table(*args, **kargs):
808    __make_table(lambda l:"%%-%is" % l, lambda l:"%%-%is" % l, "", *args, **kargs)
809
810def make_lined_table(*args, **kargs):
811    __make_table(lambda l:"%%-%is |" % l, lambda l:"%%-%is |" % l, "",
812                 seplinefunc=lambda a,x:"+".join(map(lambda y:"-"*(y+2), [a-1]+x+[-2])),
813                 *args, **kargs)
814
815def make_tex_table(*args, **kargs):
816    __make_table(lambda l: "%s", lambda l: "& %s", "\\\\", seplinefunc=lambda a,x:"\\hline", *args, **kargs)
817
818