1from __future__ import print_function
2from __future__ import absolute_import
3# Copyright (c) 2003-2016 CORE Security Technologies
4#
5# This software is provided under under a slightly modified version
6# of the Apache Software License. See the accompanying LICENSE file
7# for more information.
8#
9
10
11# -*- mode: python; tab-width: 4 -*-
12#
13# Copyright (C) 2001 Michael Teo <michaelteo@bigfoot.com>
14# nmb.py - NetBIOS library
15#
16# This software is provided 'as-is', without any express or implied warranty.
17# In no event will the author be held liable for any damages arising from the
18# use of this software.
19#
20# Permission is granted to anyone to use this software for any purpose,
21# including commercial applications, and to alter it and redistribute it
22# freely, subject to the following restrictions:
23#
24# 1. The origin of this software must not be misrepresented; you must not
25#    claim that you wrote the original software. If you use this software
26#    in a product, an acknowledgment in the product documentation would be
27#    appreciated but is not required.
28#
29# 2. Altered source versions must be plainly marked as such, and must not be
30#    misrepresented as being the original software.
31#
32# 3. This notice cannot be removed or altered from any source distribution.
33#
34# Altered source done by Alberto Solino (@agsolino)
35
36import socket
37import string
38import re
39import select
40import errno
41from random import randint
42from struct import pack, unpack
43import time
44
45from .structure import Structure
46
47CVS_REVISION = '$Revision: 526 $'
48
49# Taken from socket module reference
50INADDR_ANY = '0.0.0.0'
51BROADCAST_ADDR = '<broadcast>'
52
53# Default port for NetBIOS name service
54NETBIOS_NS_PORT = 137
55# Default port for NetBIOS session service
56NETBIOS_SESSION_PORT = 139
57
58# Default port for SMB session service
59SMB_SESSION_PORT = 445
60
61# Owner Node Type Constants
62NODE_B = 0x0000
63NODE_P = 0x2000
64NODE_M = 0x4000
65NODE_RESERVED = 0x6000
66NODE_GROUP = 0x8000
67NODE_UNIQUE = 0x0
68
69# Name Type Constants
70TYPE_UNKNOWN = 0x01
71TYPE_WORKSTATION = 0x00
72TYPE_CLIENT = 0x03
73TYPE_SERVER = 0x20
74TYPE_DOMAIN_MASTER = 0x1B
75TYPE_DOMAIN_CONTROLLER = 0x1C
76TYPE_MASTER_BROWSER = 0x1D
77TYPE_BROWSER = 0x1E
78TYPE_NETDDE  = 0x1F
79TYPE_STATUS = 0x21
80
81# Opcodes values
82OPCODE_QUERY = 0
83OPCODE_REGISTRATION = 0x5
84OPCODE_RELEASE = 0x6
85OPCODE_WACK = 0x7
86OPCODE_REFRESH = 0x8
87OPCODE_REQUEST = 0
88OPCODE_RESPONSE = 0x10
89
90# NM_FLAGS
91NM_FLAGS_BROADCAST = 0x1
92NM_FLAGS_UNICAST = 0
93NM_FLAGS_RA = 0x8
94NM_FLAGS_RD = 0x10
95NM_FLAGS_TC = 0x20
96NM_FLAGS_AA = 0x40
97
98# QUESTION_TYPE
99QUESTION_TYPE_NB = 0x20     # NetBIOS general Name Service Resource Record
100QUESTION_TYPE_NBSTAT = 0x21 # NetBIOS NODE STATUS Resource Record
101# QUESTION_CLASS
102QUESTION_CLASS_IN = 0x1     # Internet class
103
104# RR_TYPE Resource Record Type code
105RR_TYPE_A = 0x1               # IP address Resource Record
106RR_TYPE_NS = 0x2              # Name Server Resource Record
107RR_TYPE_NULL = 0xA          # NULL Resource Record
108RR_TYPE_NB = 0x20           # NetBIOS general Name Service Resource Record
109RR_TYPE_NBSTAT = 0x21       # NetBIOS NODE STATUS Resource Record
110
111# Resource Record Class
112RR_CLASS_IN = 1             # Internet class
113
114# RCODE values
115RCODE_FMT_ERR   = 0x1       # Format Error.  Request was invalidly formatted.
116RCODE_SRV_ERR   = 0x2       # Server failure.  Problem with NBNS, cannot process name.
117RCODE_IMP_ERR   = 0x4       # Unsupported request error.  Allowable only for challenging NBNS when gets an Update type
118                            # registration request.
119RCODE_RFS_ERR   = 0x5       # Refused error.  For policy reasons server will not register this name from this host.
120RCODE_ACT_ERR   = 0x6       # Active error.  Name is owned by another node.
121RCODE_CFT_ERR   = 0x7       # Name in conflict error.  A UNIQUE name is owned by more than one node.
122
123# NAME_FLAGS
124NAME_FLAGS_PRM = 0x0200       # Permanent Name Flag.  If one (1) then entry is for the permanent node name.  Flag is zero
125                            # (0) for all other names.
126NAME_FLAGS_ACT = 0x0400       # Active Name Flag.  All entries have this flag set to one (1).
127NAME_FLAG_CNF  = 0x0800       # Conflict Flag.  If one (1) then name on this node is in conflict.
128NAME_FLAG_DRG  = 0x1000       # Deregister Flag.  If one (1) then this name is in the process of being deleted.
129
130NAME_TYPES = { TYPE_UNKNOWN: 'Unknown', TYPE_WORKSTATION: 'Workstation', TYPE_CLIENT: 'Client',
131               TYPE_SERVER: 'Server', TYPE_MASTER_BROWSER: 'Master Browser', TYPE_BROWSER: 'Browser Server',
132               TYPE_DOMAIN_MASTER: 'Domain Master' , TYPE_NETDDE: 'NetDDE Server'}
133# NetBIOS Session Types
134NETBIOS_SESSION_MESSAGE = 0x0
135NETBIOS_SESSION_REQUEST = 0x81
136NETBIOS_SESSION_POSITIVE_RESPONSE = 0x82
137NETBIOS_SESSION_NEGATIVE_RESPONSE = 0x83
138NETBIOS_SESSION_RETARGET_RESPONSE = 0x84
139NETBIOS_SESSION_KEEP_ALIVE = 0x85
140
141
142def strerror(errclass, errcode):
143    if errclass == ERRCLASS_OS:
144        return 'OS Error', str(errcode)
145    elif errclass == ERRCLASS_QUERY:
146        return 'Query Error', QUERY_ERRORS.get(errcode, 'Unknown error')
147    elif errclass == ERRCLASS_SESSION:
148        return 'Session Error', SESSION_ERRORS.get(errcode, 'Unknown error')
149    else:
150        return 'Unknown Error Class', 'Unknown Error'
151
152
153
154class NetBIOSError(Exception): pass
155class NetBIOSTimeout(Exception):
156    def __init__(self, message = 'The NETBIOS connection with the remote host timed out.'):
157        Exception.__init__(self, message)
158
159class NBResourceRecord:
160    def __init__(self, data = 0):
161        self._data = data
162        try:
163            if self._data:
164                self.rr_name = (re.split('\x00',data))[0]
165                offset = len(self.rr_name)+1
166                self.rr_type  = unpack('>H', self._data[offset:offset+2])[0]
167                self.rr_class = unpack('>H', self._data[offset+2: offset+4])[0]
168                self.ttl = unpack('>L',self._data[offset+4:offset+8])[0]
169                self.rdlength = unpack('>H', self._data[offset+8:offset+10])[0]
170                self.rdata = self._data[offset+10:offset+10+self.rdlength]
171                offset = self.rdlength - 2
172                self.unit_id = data[offset:offset+6]
173            else:
174                self.rr_name = ''
175                self.rr_type = 0
176                self.rr_class = 0
177                self.ttl = 0
178                self.rdlength = 0
179                self.rdata = ''
180                self.unit_id = ''
181        except Exception:
182                raise NetBIOSError( 'Wrong packet format ' )
183
184    def set_rr_name(self, name):
185        self.rr_name = name
186    def set_rr_type(self, name):
187        self.rr_type = name
188    def set_rr_class(self,cl):
189        self.rr_class = cl
190    def set_ttl(self,ttl):
191        self.ttl = ttl
192    def set_rdata(self,rdata):
193        self.rdata = rdata
194        self.rdlength = len(rdata)
195    def get_unit_id(self):
196        return self.unit_id
197    def get_rr_name(self):
198        return self.rr_name
199    def get_rr_class(self):
200        return self.rr_class
201    def get_ttl(self):
202        return self.ttl
203    def get_rdlength(self):
204        return self.rdlength
205    def get_rdata(self):
206        return self.rdata
207    def rawData(self):
208        return self.rr_name + pack('!HHLH',self.rr_type, self.rr_class, self.ttl, self.rdlength) + self.rdata
209
210class NBNodeStatusResponse(NBResourceRecord):
211    def __init__(self, data = 0):
212        NBResourceRecord.__init__(self,data)
213        self.num_names = 0
214        self.node_names = [ ]
215        self.statstics = ''
216        self.mac = '00-00-00-00-00-00'
217        try:
218            if data:
219                self._data = self.get_rdata()
220                self.num_names = unpack('>B',self._data[:1])[0]
221                offset = 1
222                for i in range(0, self.num_names):
223                    name = self._data[offset:offset + 15]
224                    type,flags = unpack('>BH', self._data[offset + 15: offset + 18])
225                    offset += 18
226                    self.node_names.append(NBNodeEntry(name, type ,flags))
227                self.set_mac_in_hexa(self.get_unit_id())
228        except Exception:
229            raise NetBIOSError( 'Wrong packet format ' )
230
231    def set_mac_in_hexa(self, data):
232        data_aux = ''
233        for d in data:
234            if data_aux == '':
235                data_aux = '%02x' % ord(d)
236            else:
237                data_aux += '-%02x' % ord(d)
238        self.mac = string.upper(data_aux)
239
240    def get_num_names(self):
241        return self.num_names
242    def get_mac(self):
243        return self.mac
244    def set_num_names(self, num):
245        self.num_names = num
246    def get_node_names(self):
247        return self.node_names
248    def add_node_name(self,node_names):
249        self.node_names.append(node_names)
250        self.num_names += 1
251    def rawData(self):
252        res = pack('!B', self.num_names )
253        for i in range(0, self.num_names):
254            res += self.node_names[i].rawData()
255
256class NBPositiveNameQueryResponse(NBResourceRecord):
257    def __init__(self, data = 0):
258        NBResourceRecord.__init__(self, data)
259        self.addr_entries = [ ]
260        if data:
261                self._data = self.get_rdata()
262                _qn_length, qn_name, qn_scope = decode_name(data)
263                self._netbios_name = string.rstrip(qn_name[:-1]) + qn_scope
264                self._name_type = ord(qn_name[-1])
265                self._nb_flags = unpack('!H', self._data[:2])
266                offset = 2
267                while offset<len(self._data):
268                    self.addr_entries.append('%d.%d.%d.%d' % unpack('4B', (self._data[offset:offset+4])))
269                    offset += 4
270
271    def get_netbios_name(self):
272        return self._netbios_name
273
274    def get_name_type(self):
275        return self._name_type
276
277    def get_addr_entries(self):
278        return self.addr_entries
279
280class NetBIOSPacket:
281    """ This is a packet as defined in RFC 1002 """
282    def __init__(self, data = 0):
283        self.name_trn_id = 0x0  # Transaction ID for Name Service Transaction.
284                                #   Requestor places a unique value for each active
285                                #   transaction.  Responder puts NAME_TRN_ID value
286                                #   from request packet in response packet.
287        self.opcode = 0         # Packet type code
288        self.nm_flags = 0       # Flags for operation
289        self.rcode = 0          # Result codes of request.
290        self.qdcount = 0        # Unsigned 16 bit integer specifying the number of entries in the question section of a Name
291        self.ancount = 0        # Unsigned 16 bit integer specifying the number of
292                                # resource records in the answer section of a Name
293                                # Service packet.
294        self.nscount = 0        # Unsigned 16 bit integer specifying the number of
295                                # resource records in the authority section of a
296                                # Name Service packet.
297        self.arcount = 0        # Unsigned 16 bit integer specifying the number of
298                                # resource records in the additional records
299                                # section of a Name Service packeT.
300        self.questions = ''
301        self.answers = ''
302        if data == 0:
303            self._data = ''
304        else:
305            try:
306                self._data = data
307                self.opcode = ord(data[2]) >> 3
308                self.nm_flags = ((ord(data[2]) & 0x3) << 4) | ((ord(data[3]) & 0xf0) >> 4)
309                self.name_trn_id = unpack('>H', self._data[:2])[0]
310                self.rcode = ord(data[3]) & 0x0f
311                self.qdcount = unpack('>H', self._data[4:6])[0]
312                self.ancount = unpack('>H', self._data[6:8])[0]
313                self.nscount = unpack('>H', self._data[8:10])[0]
314                self.arcount = unpack('>H', self._data[10:12])[0]
315                self.answers = self._data[12:]
316            except Exception:
317                raise NetBIOSError( 'Wrong packet format ' )
318
319    def set_opcode(self, opcode):
320        self.opcode = opcode
321    def set_trn_id(self, trn):
322        self.name_trn_id = trn
323    def set_nm_flags(self, nm_flags):
324        self.nm_flags = nm_flags
325    def set_rcode(self, rcode):
326        self.rcode = rcode
327    def addQuestion(self, question, qtype, qclass):
328        self.qdcount += 1
329        self.questions += question + pack('!HH',qtype,qclass)
330    def get_trn_id(self):
331        return self.name_trn_id
332    def get_rcode(self):
333        return self.rcode
334    def get_nm_flags(self):
335        return self.nm_flags
336    def get_opcode(self):
337        return self.opcode
338    def get_qdcount(self):
339        return self.qdcount
340    def get_ancount(self):
341        return self.ancount
342    def get_nscount(self):
343        return self.nscount
344    def get_arcount(self):
345        return self.arcount
346    def rawData(self):
347        secondWord = self.opcode << 11
348        secondWord |= self.nm_flags << 4
349        secondWord |= self.rcode
350        data = pack('!HHHHHH', self.name_trn_id, secondWord , self.qdcount, self.ancount, self.nscount, self.arcount) + self.questions + self.answers
351        return data
352    def get_answers(self):
353        return self.answers
354
355class NBHostEntry:
356
357    def __init__(self, nbname, nametype, ip):
358        self.__nbname = nbname
359        self.__nametype = nametype
360        self.__ip = ip
361
362    def get_nbname(self):
363        return self.__nbname
364
365    def get_nametype(self):
366        return self.__nametype
367
368    def get_ip(self):
369        return self.__ip
370
371    def __repr__(self):
372        return '<NBHostEntry instance: NBname="' + self.__nbname + '", IP="' + self.__ip + '">'
373
374class NBNodeEntry:
375
376    def __init__(self, nbname, nametype, flags):
377        self.__nbname = string.ljust(nbname,17)
378        self.__nametype = nametype
379        self.__flags = flags
380        self.__isgroup = flags & 0x8000
381        self.__nodetype = flags & 0x6000
382        self.__deleting = flags & 0x1000
383        self.__isconflict = flags & 0x0800
384        self.__isactive = flags & 0x0400
385        self.__ispermanent = flags & 0x0200
386
387    def get_nbname(self):
388        return self.__nbname
389
390    def get_nametype(self):
391        return self.__nametype
392
393    def is_group(self):
394        return self.__isgroup
395
396    def get_nodetype(self):
397        return self.__nodetype
398
399    def is_deleting(self):
400        return self.__deleting
401
402    def is_conflict(self):
403        return self.__isconflict
404
405    def is_active(self):
406        return self.__isactive
407
408    def is_permanent(self):
409        return self.__ispermanent
410
411    def set_nbname(self, name):
412        self.__nbname = string.ljust(name,17)
413
414    def set_nametype(self, type):
415        self.__nametype = type
416
417    def set_flags(self,flags):
418        self.__flags = flags
419
420    def __repr__(self):
421        s = '<NBNodeEntry instance: NBname="' + self.__nbname + '" NameType="' + NAME_TYPES[self.__nametype] + '"'
422        if self.__isactive:
423            s += ' ACTIVE'
424        if self.__isgroup:
425            s += ' GROUP'
426        if self.__isconflict:
427            s += ' CONFLICT'
428        if self.__deleting:
429            s += ' DELETING'
430        return s
431    def rawData(self):
432        return self.__nbname + pack('!BH',self.__nametype, self.__flags)
433
434
435class NetBIOS:
436
437    # Creates a NetBIOS instance without specifying any default NetBIOS domain nameserver.
438    # All queries will be sent through the servport.
439    def __init__(self, servport = NETBIOS_NS_PORT):
440        self.__servport = NETBIOS_NS_PORT
441        self.__nameserver = None
442        self.__broadcastaddr = BROADCAST_ADDR
443        self.mac = '00-00-00-00-00-00'
444
445    def _setup_connection(self, dstaddr):
446        port = randint(10000, 60000)
447        af, socktype, proto, _canonname, _sa = socket.getaddrinfo(dstaddr, port, socket.AF_INET, socket.SOCK_DGRAM)[0]
448        s = socket.socket(af, socktype, proto)
449        has_bind = 1
450        for _i in range(0, 10):
451        # We try to bind to a port for 10 tries
452            try:
453                s.bind(( INADDR_ANY, randint(10000, 60000) ))
454                s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
455                has_bind = 1
456            except socket.error:
457                pass
458        if not has_bind:
459            raise NetBIOSError( 'Cannot bind to a good UDP port', ERRCLASS_OS, errno.EAGAIN)
460        self.__sock = s
461
462    # Set the default NetBIOS domain nameserver.
463    def set_nameserver(self, nameserver):
464        self.__nameserver = nameserver
465
466    # Return the default NetBIOS domain nameserver, or None if none is specified.
467    def get_nameserver(self):
468        return self.__nameserver
469
470    # Set the broadcast address to be used for query.
471    def set_broadcastaddr(self, broadcastaddr):
472        self.__broadcastaddr = broadcastaddr
473
474    # Return the broadcast address to be used, or BROADCAST_ADDR if default broadcast address is used.
475    def get_broadcastaddr(self):
476        return self.__broadcastaddr
477
478    # Returns a NBPositiveNameQueryResponse instance containing the host information for nbname.
479    # If a NetBIOS domain nameserver has been specified, it will be used for the query.
480    # Otherwise, the query is broadcasted on the broadcast address.
481    def gethostbyname(self, nbname, qtype = TYPE_WORKSTATION, scope = None, timeout = 1):
482        return self.__queryname(nbname, self.__nameserver, qtype, scope, timeout)
483
484    # Returns a list of NBNodeEntry instances containing node status information for nbname.
485    # If destaddr contains an IP address, then this will become an unicast query on the destaddr.
486    # Raises NetBIOSTimeout if timeout (in secs) is reached.
487    # Raises NetBIOSError for other errors
488    def getnodestatus(self, nbname, destaddr = None, type = TYPE_WORKSTATION, scope = None, timeout = 1):
489        if destaddr:
490            return self.__querynodestatus(nbname, destaddr, type, scope, timeout)
491        else:
492            return self.__querynodestatus(nbname, self.__nameserver, type, scope, timeout)
493
494    def getnetbiosname(self, ip):
495        entries = self.getnodestatus('*',ip)
496        entries = filter(lambda x:x.get_nametype() == TYPE_SERVER, entries)
497        return entries[0].get_nbname().strip()
498
499    def getmacaddress(self):
500        return self.mac
501
502    def __queryname(self, nbname, destaddr, qtype, scope, timeout, retries = 0):
503        self._setup_connection(destaddr)
504        trn_id = randint(1, 32000)
505        p = NetBIOSPacket()
506        p.set_trn_id(trn_id)
507        netbios_name = nbname.upper()
508        qn_label = encode_name(netbios_name, qtype, scope)
509        p.addQuestion(qn_label, QUESTION_TYPE_NB, QUESTION_CLASS_IN)
510        p.set_nm_flags(NM_FLAGS_RD)
511        if not destaddr:
512            p.set_nm_flags(p.get_nm_flags() | NM_FLAGS_BROADCAST)
513            destaddr = self.__broadcastaddr
514        req = p.rawData()
515
516        tries = retries
517        while 1:
518            self.__sock.sendto(req, ( destaddr, self.__servport ))
519            try:
520                ready, _, _ = select.select([ self.__sock.fileno() ], [ ] , [ ], timeout)
521                if not ready:
522                    if tries:
523                        # Retry again until tries == 0
524                        tries -= 1
525                    else:
526                        raise NetBIOSTimeout
527                else:
528                    data, _ = self.__sock.recvfrom(65536, 0)
529
530                    res = NetBIOSPacket(data)
531                    if res.get_trn_id() == p.get_trn_id():
532                        if res.get_rcode():
533                            if res.get_rcode() == 0x03:
534                                return None
535                            else:
536                                raise NetBIOSError( 'Negative name query response', ERRCLASS_QUERY, res.get_rcode())
537
538                        if res.get_ancount() != 1:
539                            raise NetBIOSError( 'Malformed response')
540
541                        return NBPositiveNameQueryResponse(res.get_answers())
542            except select.error as ex:
543                if ex[0] != errno.EINTR and ex[0] != errno.EAGAIN:
544                    raise NetBIOSError( 'Error occurs while waiting for response', ERRCLASS_OS, ex[0])
545                raise
546
547
548    def __querynodestatus(self, nbname, destaddr, type, scope, timeout):
549        self._setup_connection(destaddr)
550        trn_id = randint(1, 32000)
551        p = NetBIOSPacket()
552        p.set_trn_id(trn_id)
553        netbios_name = string.upper(nbname)
554        qn_label = encode_name(netbios_name, type, scope)
555        p.addQuestion(qn_label, QUESTION_TYPE_NBSTAT, QUESTION_CLASS_IN)
556
557        if not destaddr:
558            p.set_nm_flags(NM_FLAGS_BROADCAST)
559            destaddr = self.__broadcastaddr
560        req = p.rawData()
561        tries = 3
562        while 1:
563            try:
564                self.__sock.sendto(req, 0, ( destaddr, self.__servport ))
565                ready, _, _ = select.select([ self.__sock.fileno() ], [ ] , [ ], timeout)
566                if not ready:
567                    if tries:
568                        # Retry again until tries == 0
569                        tries -= 1
570                    else:
571                        raise NetBIOSTimeout
572                else:
573                    try:
574                        data, _ = self.__sock.recvfrom(65536, 0)
575                    except Exception as e:
576                        raise NetBIOSError("recvfrom error: %s" % str(e))
577                    self.__sock.close()
578                    res = NetBIOSPacket(data)
579                    if res.get_trn_id() == p.get_trn_id():
580                        if res.get_rcode():
581                            if res.get_rcode() == 0x03:
582                                # I'm just guessing here
583                                raise NetBIOSError("Cannot get data from server")
584                            else:
585                                raise NetBIOSError( 'Negative name query response', ERRCLASS_QUERY, res.get_rcode())
586                        answ = NBNodeStatusResponse(res.get_answers())
587                        self.mac = answ.get_mac()
588                        return answ.get_node_names()
589            except select.error as ex:
590                if ex[0] != errno.EINTR and ex[0] != errno.EAGAIN:
591                    raise NetBIOSError( 'Error occurs while waiting for response', ERRCLASS_OS, ex[0])
592            except socket.error as ex:
593                raise NetBIOSError('Connection error: %s' % str(ex))
594
595# Perform first and second level encoding of name as specified in RFC 1001 (Section 4)
596def encode_name(name, type, scope):
597    if name == '*':
598        name += '\0' * 15
599    elif len(name) > 15:
600        name = name[:15] + chr(type)
601    else:
602        name = string.ljust(name, 15) + chr(type)
603
604    encoded_name = chr(len(name) * 2) + re.sub('.', _do_first_level_encoding, name)
605    if scope:
606        encoded_scope = ''
607        for s in string.split(scope, '.'):
608            encoded_scope = encoded_scope + chr(len(s)) + s
609        return encoded_name + encoded_scope + '\0'
610    else:
611        return encoded_name + '\0'
612
613# Internal method for use in encode_name()
614def _do_first_level_encoding(m):
615    s = ord(m.group(0))
616    return string.uppercase[s >> 4] + string.uppercase[s & 0x0f]
617
618def decode_name(name):
619    name_length = ord(name[0])
620    assert name_length == 32
621
622    decoded_name = re.sub('..', _do_first_level_decoding, name[1:33])
623    if name[33] == '\0':
624        return 34, decoded_name, ''
625    else:
626        decoded_domain = ''
627        offset = 34
628        while 1:
629            domain_length = ord(name[offset])
630            if domain_length == 0:
631                break
632            decoded_domain = '.' + name[offset:offset + domain_length]
633            offset += domain_length
634        return offset + 1, decoded_name, decoded_domain
635
636def _do_first_level_decoding(m):
637    s = m.group(0)
638    return chr(((ord(s[0]) - ord('A')) << 4) | (ord(s[1]) - ord('A')))
639
640
641
642class NetBIOSSessionPacket:
643    def __init__(self, data = 0):
644        self.type = 0x0
645        self.flags = 0x0
646        self.length = 0x0
647        if data == 0:
648            self._trailer = ''
649        else:
650            try:
651                self.type = ord(data[0])
652                if self.type == NETBIOS_SESSION_MESSAGE:
653                    self.length = ord(data[1]) << 16 | (unpack('!H', data[2:4])[0])
654                else:
655                    self.flags = ord(data[1])
656                    self.length = unpack('!H', data[2:4])[0]
657
658                self._trailer = data[4:]
659            except:
660                raise NetBIOSError( 'Wrong packet format ' )
661
662    def set_type(self, type):
663        self.type = type
664    def get_type(self):
665        return self.type
666    def rawData(self):
667        if self.type == NETBIOS_SESSION_MESSAGE:
668            data = pack('!BBH',self.type,self.length >> 16,self.length & 0xFFFF) + self._trailer
669        else:
670            data = pack('!BBH',self.type,self.flags,self.length) + self._trailer
671        return data
672    def set_trailer(self,data):
673        self._trailer = data
674        self.length = len(data)
675    def get_length(self):
676        return self.length
677    def get_trailer(self):
678        return self._trailer
679
680class NetBIOSSession:
681    def __init__(self, myname, remote_name, remote_host, remote_type = TYPE_SERVER, sess_port = NETBIOS_SESSION_PORT, timeout = None, local_type = TYPE_WORKSTATION, sock = None):
682        if len(myname) > 15:
683            self.__myname = string.upper(myname[:15])
684        else:
685            self.__myname = string.upper(myname)
686        self.__local_type = local_type
687
688        assert remote_name
689        # if destination port SMB_SESSION_PORT and remote name *SMBSERVER, we're changing it to its IP address
690        # helping solving the client mistake ;)
691        if remote_name == '*SMBSERVER' and sess_port == SMB_SESSION_PORT:
692            remote_name = remote_host
693        # If remote name is *SMBSERVER let's try to query its name.. if can't be guessed, continue and hope for the best
694        if remote_name == '*SMBSERVER':
695            nb = NetBIOS()
696
697            try:
698                res = nb.getnetbiosname(remote_host)
699            except:
700                res = None
701                pass
702
703            if res is not None:
704                remote_name = res
705
706        if len(remote_name) > 15:
707            self.__remote_name = string.upper(remote_name[:15])
708        else:
709            self.__remote_name = string.upper(remote_name)
710        self.__remote_type = remote_type
711
712        self.__remote_host = remote_host
713
714        if sock is not None:
715            # We are acting as a server
716            self._sock = sock
717        else:
718            self._sock = self._setup_connection((remote_host, sess_port))
719
720        if sess_port == NETBIOS_SESSION_PORT:
721            self._request_session(remote_type, local_type, timeout)
722
723    def get_myname(self):
724        return self.__myname
725
726    def get_mytype(self):
727        return self.__local_type
728
729    def get_remote_host(self):
730        return self.__remote_host
731
732    def get_remote_name(self):
733        return self.__remote_name
734
735    def get_remote_type(self):
736        return self.__remote_type
737
738    def close(self):
739        self._sock.close()
740
741    def get_socket(self):
742        return self._sock
743
744class NetBIOSUDPSessionPacket(Structure):
745    TYPE_DIRECT_UNIQUE = 16
746    TYPE_DIRECT_GROUP  = 17
747
748    FLAGS_MORE_FRAGMENTS = 1
749    FLAGS_FIRST_FRAGMENT = 2
750    FLAGS_B_NODE         = 0
751
752    structure = (
753        ('Type','B=16'),    # Direct Unique Datagram
754        ('Flags','B=2'),    # FLAGS_FIRST_FRAGMENT
755        ('ID','<H'),
756        ('_SourceIP','>L'),
757        ('SourceIP','"'),
758        ('SourcePort','>H=138'),
759        ('DataLegth','>H-Data'),
760        ('Offset','>H=0'),
761        ('SourceName','z'),
762        ('DestinationName','z'),
763        ('Data',':'),
764    )
765
766    def getData(self):
767        addr = self['SourceIP'].split('.')
768        addr = [int(x) for x in addr]
769        addr = (((addr[0] << 8) + addr[1] << 8) + addr[2] << 8) + addr[3]
770        self['_SourceIP'] = addr
771        return Structure.getData(self)
772
773    def get_trailer(self):
774        return self['Data']
775
776class NetBIOSUDPSession(NetBIOSSession):
777    def _setup_connection(self, peer):
778        af, socktype, proto, canonname, sa = socket.getaddrinfo(peer[0], peer[1], 0, socket.SOCK_DGRAM)[0]
779        sock = socket.socket(af, socktype, proto)
780        sock.connect(sa)
781
782        sock = socket.socket(af, socktype, proto)
783        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
784        sock.bind((INADDR_ANY, 138))
785        self.peer = peer
786        return sock
787
788    def _request_session(self, remote_type, local_type, timeout = None):
789        pass
790
791    def next_id(self):
792        if hasattr(self, '__dgram_id'):
793            answer = self.__dgram_id
794        else:
795            self.__dgram_id = randint(1,65535)
796            answer = self.__dgram_id
797        self.__dgram_id += 1
798        return answer
799
800    def send_packet(self, data):
801        # Yes... I know...
802        self._sock.connect(self.peer)
803
804        p = NetBIOSUDPSessionPacket()
805        p['ID'] = self.next_id()
806        p['SourceIP'] = self._sock.getsockname()[0]
807        p['SourceName'] = encode_name(self.get_myname(), self.get_mytype(), '')[:-1]
808        p['DestinationName'] = encode_name(self.get_remote_name(), self.get_remote_type(), '')[:-1]
809        p['Data'] = data
810
811        self._sock.sendto(str(p), self.peer)
812        self._sock.close()
813
814        self._sock = self._setup_connection(self.peer)
815
816    def recv_packet(self, timeout = None):
817        # The next loop is a workaround for a bigger problem:
818        # When data reaches higher layers, the lower headers are lost,
819        # and with them, for example, the source IP. Hence, SMB users
820        # can't know where packets are comming from... we need a better
821        # solution, right now, we will filter everything except packets
822        # coming from the remote_host specified in __init__()
823
824        while 1:
825            data, peer = self._sock.recvfrom(8192)
826#            print "peer: %r  self.peer: %r" % (peer, self.peer)
827            if peer == self.peer: break
828
829        return NetBIOSUDPSessionPacket(data)
830
831class NetBIOSTCPSession(NetBIOSSession):
832    def __init__(self, myname, remote_name, remote_host, remote_type = TYPE_SERVER, sess_port = NETBIOS_SESSION_PORT, timeout = None, local_type = TYPE_WORKSTATION, sock = None, select_poll = False):
833        self.__select_poll = select_poll
834        if self.__select_poll:
835            self.read_function = self.polling_read
836        else:
837            self.read_function = self.non_polling_read
838        NetBIOSSession.__init__(self, myname, remote_name, remote_host, remote_type = remote_type, sess_port = sess_port, timeout = timeout, local_type = local_type, sock=sock)
839
840
841    def _setup_connection(self, peer):
842        try:
843            af, socktype, proto, canonname, sa = socket.getaddrinfo(peer[0], peer[1], 0, socket.SOCK_STREAM)[0]
844            sock = socket.socket(af, socktype, proto)
845            sock.connect(sa)
846        except socket.error as e:
847            raise socket.error("Connection error (%s:%s)" % (peer[0], peer[1]), e)
848        return sock
849
850    def send_packet(self, data):
851        p = NetBIOSSessionPacket()
852        p.set_type(NETBIOS_SESSION_MESSAGE)
853        p.set_trailer(data)
854        self._sock.send(p.rawData())
855
856    def recv_packet(self, timeout = None):
857        data = self.__read(timeout)
858        return NetBIOSSessionPacket(data)
859
860    def _request_session(self, remote_type, local_type, timeout = None):
861        p = NetBIOSSessionPacket()
862        remote_name = encode_name(self.get_remote_name(), remote_type, '')
863        myname = encode_name(self.get_myname(), local_type, '')
864        p.set_type(NETBIOS_SESSION_REQUEST)
865        p.set_trailer(remote_name + myname)
866
867        self._sock.send(p.rawData())
868        while 1:
869            p = self.recv_packet(timeout)
870            if p.get_type() == NETBIOS_SESSION_NEGATIVE_RESPONSE:
871                raise NetBIOSError( 'Cannot request session', ERRCLASS_SESSION, ord(p.get_trailer()[0]))
872            elif p.get_type() == NETBIOS_SESSION_POSITIVE_RESPONSE:
873                break
874            else:
875                # Ignore all other messages, most probably keepalive messages
876                pass
877
878    def polling_read(self, read_length, timeout):
879        data = ''
880        if timeout is None:
881            timeout = 3600
882
883        time_left = timeout
884        CHUNK_TIME = 0.025
885        bytes_left = read_length
886
887        while bytes_left > 0:
888            try:
889                ready, _, _ = select.select([self._sock.fileno() ], [ ], [ ], 0)
890
891                if not ready:
892                    if time_left <= 0:
893                        raise NetBIOSTimeout
894                    else:
895                        time.sleep(CHUNK_TIME)
896                        time_left -= CHUNK_TIME
897                        continue
898
899                received = self._sock.recv(bytes_left)
900                if len(received) == 0:
901                    raise NetBIOSError( 'Error while reading from remote', ERRCLASS_OS, None)
902
903                data = data + received
904                bytes_left = read_length - len(data)
905            except select.error as ex:
906                if ex[0] != errno.EINTR and ex[0] != errno.EAGAIN:
907                    raise NetBIOSError( 'Error occurs while reading from remote', ERRCLASS_OS, ex[0])
908
909        return data
910
911    def non_polling_read(self, read_length, timeout):
912        data = ''
913        bytes_left = read_length
914
915        while bytes_left > 0:
916            try:
917                ready, _, _ = select.select([self._sock.fileno() ], [ ], [ ], timeout)
918
919                if not ready:
920                        raise NetBIOSTimeout
921
922                received = self._sock.recv(bytes_left)
923                if len(received) == 0:
924                    raise NetBIOSError( 'Error while reading from remote', ERRCLASS_OS, None)
925
926                data = data + received
927                bytes_left = read_length - len(data)
928            except select.error as ex:
929                if ex[0] != errno.EINTR and ex[0] != errno.EAGAIN:
930                    raise NetBIOSError( 'Error occurs while reading from remote', ERRCLASS_OS, ex[0])
931
932        return data
933
934    def __read(self, timeout = None):
935        data = self.read_function(4, timeout)
936        type, flags, length = unpack('>ccH', data)
937        if ord(type) == NETBIOS_SESSION_MESSAGE:
938            length |= ord(flags) << 16
939        else:
940            if ord(flags) & 0x01:
941                length |= 0x10000
942        data2 = self.read_function(length, timeout)
943
944        return data + data2
945
946ERRCLASS_QUERY = 0x00
947ERRCLASS_SESSION = 0xf0
948ERRCLASS_OS = 0xff
949
950QUERY_ERRORS = { 0x01: 'Request format error. Please file a bug report.',
951                 0x02: 'Internal server error',
952                 0x03: 'Name does not exist',
953                 0x04: 'Unsupported request',
954                 0x05: 'Request refused'
955                 }
956
957SESSION_ERRORS = { 0x80: 'Not listening on called name',
958                   0x81: 'Not listening for calling name',
959                   0x82: 'Called name not present',
960                   0x83: 'Sufficient resources',
961                   0x8f: 'Unspecified error'
962                   }
963
964def main():
965    def get_netbios_host_by_name(name):
966        n = NetBIOS()
967        n.set_broadcastaddr('255.255.255.255') # To avoid use "<broadcast>" in socket
968        for qtype in (TYPE_WORKSTATION, TYPE_CLIENT, TYPE_SERVER, TYPE_DOMAIN_MASTER, TYPE_DOMAIN_CONTROLLER):
969            try:
970                addrs = n.gethostbyname(name, qtype = qtype).get_addr_entries()
971            except NetBIOSTimeout:
972                continue
973            else:
974                return addrs
975        raise Exception("Host not found")
976
977
978    n = get_netbios_host_by_name("some-host")
979    print(n)
980
981if __name__ == '__main__':
982    main()
983