source: scapy/utils.py @ 1229:33f69be3ddcd

Revision 1229:33f69be3ddcd, 21.8 KB checked in by Dirk Loss <mail@…>, 3 years ago (diff)

Added (simple) module-level docstrings for all modules

Line 
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
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   
564
565class PcapReader(RawPcapReader):
566    def __init__(self, filename):
567        RawPcapReader.__init__(self, filename)
568        try:
569            self.LLcls = conf.l2types[self.linktype]
570        except KeyError:
571            warning("PcapReader: unknown LL type [%i]/[%#x]. Using Raw packets" % (self.linktype,self.linktype))
572            self.LLcls = conf.raw_layer
573    def read_packet(self, size=MTU):
574        rp = RawPcapReader.read_packet(self,size)
575        if rp is None:
576            return None
577        s,(sec,usec,wirelen) = rp
578       
579        try:
580            p = self.LLcls(s)
581        except KeyboardInterrupt:
582            raise
583        except:
584            if conf.debug_dissector:
585                raise
586            p = conf.raw_layer(s)
587        p.time = sec+0.000001*usec
588        return p
589    def read_all(self,count=-1):
590        res = RawPcapReader.read_all(self, count)
591        import plist
592        return plist.PacketList(res,name = os.path.basename(self.filename))
593    def recv(self, size=MTU):
594        return self.read_packet(size)
595       
596
597
598class RawPcapWriter:
599    """A stream PCAP writer with more control than wrpcap()"""
600    def __init__(self, filename, linktype=None, gz=False, endianness="", append=False, sync=False):
601        """
602        linktype: force linktype to a given value. If None, linktype is taken
603                  from the first writter packet
604        gz: compress the capture on the fly
605        endianness: force an endianness (little:"<", big:">"). Default is native
606        append: append packets to the capture file instead of truncating it
607        sync: do not bufferize writes to the capture file
608        """
609       
610        self.linktype = linktype
611        self.header_present = 0
612        self.append=append
613        self.gz = gz
614        self.endian = endianness
615        self.filename=filename
616        self.sync=sync
617        bufsz=4096
618        if sync:
619            bufsz=0
620
621        self.f = [open,gzip.open][gz](filename,append and "ab" or "wb", gz and 9 or bufsz)
622       
623    def fileno(self):
624        return self.f.fileno()
625
626    def _write_header(self, pkt):
627        self.header_present=1
628
629        if self.append:
630            # Even if prone to race conditions, this seems to be
631            # safest way to tell whether the header is already present
632            # because we have to handle compressed streams that
633            # are not as flexible as basic files
634            g = [open,gzip.open][self.gz](self.filename,"rb")
635            if g.read(16):
636                return
637           
638        self.f.write(struct.pack(self.endian+"IHHIIII", 0xa1b2c3d4L,
639                                 2, 4, 0, 0, MTU, self.linktype))
640        self.f.flush()
641   
642
643    def write(self, pkt):
644        """accepts a either a single packet or a list of packets
645        to be written to the dumpfile
646        """
647        if not self.header_present:
648            self._write_header(pkt)
649        if type(pkt) is str:
650            self._write_packet(pkt)
651        else:
652            for p in pkt:
653                self._write_packet(p)
654
655    def _write_packet(self, packet, sec=None, usec=None, caplen=None, wirelen=None):
656        """writes a single packet to the pcap file
657        """
658        if caplen is None:
659            caplen = len(packet)
660        if wirelen is None:
661            wirelen = caplen
662        if sec is None or usec is None:
663            t=time.time()
664            it = int(t)
665            if sec is None:
666                sec = it
667            if usec is None:
668                usec = int(round((t-it)*1000000))
669        self.f.write(struct.pack(self.endian+"IIII", sec, usec, caplen, wirelen))
670        self.f.write(packet)
671        if self.gz and self.sync:
672            self.f.flush()
673
674    def flush(self):
675        return self.f.flush()
676    def close(self):
677        return self.f.close()
678               
679class PcapWriter(RawPcapWriter):
680    def _write_header(self, pkt):
681        if self.linktype == None:
682            if type(pkt) is list or type(pkt) is tuple or isinstance(pkt,BasePacketList):
683                pkt = pkt[0]
684            try:
685                self.linktype = conf.l2types[pkt.__class__]
686            except KeyError:
687                warning("PcapWriter: unknown LL type for %s. Using type 1 (Ethernet)" % pkt.__class__.__name__)
688                self.linktype = 1
689        RawPcapWriter._write_header(self, pkt)
690
691    def _write_packet(self, packet):       
692        sec = int(packet.time)
693        usec = int(round((packet.time-sec)*1000000))
694        s = str(packet)
695        caplen = len(s)
696        RawPcapWriter._write_packet(self, s, sec, usec, caplen, caplen)
697
698
699re_extract_hexcap = re.compile("^((0x)?[0-9a-fA-F]{2,}[ :\t]{,3}|) *(([0-9a-fA-F]{2} {,2}){,16})")
700
701def import_hexcap():
702    p = ""
703    try:
704        while 1:
705            l = raw_input().strip()
706            try:
707                p += re_extract_hexcap.match(l).groups()[2]
708            except:
709                warning("Parsing error during hexcap")
710                continue
711    except EOFError:
712        pass
713   
714    p = p.replace(" ","")
715    return p.decode("hex")
716       
717
718
719@conf.commands.register
720def wireshark(pktlist):
721    """Run wireshark on a list of packets"""
722    f = get_temp_file()
723    wrpcap(f, pktlist)
724    subprocess.Popen([conf.prog.wireshark, "-r", f])
725
726@conf.commands.register
727def hexedit(x):
728    x = str(x)
729    f = get_temp_file()
730    open(f,"w").write(x)
731    subprocess.call([conf.prog.hexedit, f])
732    x = open(f).read()
733    os.unlink(f)
734    return x
735
736def __make_table(yfmtfunc, fmtfunc, endline, list, fxyz, sortx=None, sorty=None, seplinefunc=None):
737    vx = {}
738    vy = {}
739    vz = {}
740    vxf = {}
741    vyf = {}
742    l = 0
743    for e in list:
744        xx,yy,zz = map(str, fxyz(e))
745        l = max(len(yy),l)
746        vx[xx] = max(vx.get(xx,0), len(xx), len(zz))
747        vy[yy] = None
748        vz[(xx,yy)] = zz
749
750    vxk = vx.keys()
751    vyk = vy.keys()
752    if sortx:
753        vxk.sort(sortx)
754    else:
755        try:
756            vxk.sort(lambda x,y:int(x)-int(y))
757        except:
758            try:
759                vxk.sort(lambda x,y: cmp(atol(x),atol(y)))
760            except:
761                vxk.sort()
762    if sorty:
763        vyk.sort(sorty)
764    else:
765        try:
766            vyk.sort(lambda x,y:int(x)-int(y))
767        except:
768            try:
769                vyk.sort(lambda x,y: cmp(atol(x),atol(y)))
770            except:
771                vyk.sort()
772
773
774    if seplinefunc:
775        sepline = seplinefunc(l, map(lambda x:vx[x],vxk))
776        print sepline
777
778    fmt = yfmtfunc(l)
779    print fmt % "",
780    for x in vxk:
781        vxf[x] = fmtfunc(vx[x])
782        print vxf[x] % x,
783    print endline
784    if seplinefunc:
785        print sepline
786    for y in vyk:
787        print fmt % y,
788        for x in vxk:
789            print vxf[x] % vz.get((x,y), "-"),
790        print endline
791    if seplinefunc:
792        print sepline
793
794def make_table(*args, **kargs):
795    __make_table(lambda l:"%%-%is" % l, lambda l:"%%-%is" % l, "", *args, **kargs)
796   
797def make_lined_table(*args, **kargs):
798    __make_table(lambda l:"%%-%is |" % l, lambda l:"%%-%is |" % l, "",
799                 seplinefunc=lambda a,x:"+".join(map(lambda y:"-"*(y+2), [a-1]+x+[-2])),
800                 *args, **kargs)
801
802def make_tex_table(*args, **kargs):
803    __make_table(lambda l: "%s", lambda l: "& %s", "\\\\", seplinefunc=lambda a,x:"\\hline", *args, **kargs)
804
Note: See TracBrowser for help on using the repository browser.