tests/python_dependencies/impacket/nmb.py

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