#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Copyright (C) 2010 Stefan J. Betz <info@stefan-betz.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os

from re import compile,M
from subprocess import Popen,PIPE
from copy import copy
from socket import gethostbyname, gethostbyaddr

# Common regular expressions
masterRe = compile("\n((?:\d{1,3}\.){3}\d{1,3})")
lookupRe = compile("\n\t([^\s]*)\s*(<(?:20|1e)>)")
sharere = compile("^\t(.*)Disk *(.*)$", M)

# Common used functions
def which(name):
    """
    Detect if name is in $PATH.
    """
    for path in os.path.expandvars('$PATH').split(':'):
        if os.access('%s/%s' % (path, name), os.F_OK):
            return '%s/%s' % (path, name)
    print 'Warning %s not found in $PATH' % name
    raise SystemExit

# Common commands
workgroupcommand = [ which('nmblookup') ]
lookupcommand = [ which('nmblookup'), '-A' ]
mastercommand = [ which('nmblookup'), '-M', '--', '-' ]
sharelookupcommand = [ which('smbclient'), '-N', '-L' ]

# Common used lookup query
def lookuphost(host,master=False):
    """
    Lookup Workgroup and NetBIOS Name of host
    """
    command = copy(lookupcommand)
    command.append(host)
    process = Popen(command,executable=command[0],stdout=PIPE,stderr=PIPE)
    process.wait()
    stdout = process.stdout.read()
    process.stdout.close()
    match = lookupRe.search(stdout)
    while match:
        if match.group(2) == "<1e>" and master:
            yield match.group(1)
        if match.group(2) == "<20>" and not master:
            yield match.group(1)
        match = lookupRe.search(stdout,match.end())

# Query Master Browsers
masters = []
process = Popen(mastercommand, stdout=PIPE,stderr=PIPE)
process.wait()
stdout = process.stdout.read()
process.stdout.close()
match = masterRe.search(stdout)
while match:
    masters.append(match.group(1))
    match = masterRe.search(stdout, match.end())
print 'Avaliable master browsers: %s' % masters

# Lookup Workgroup and Host Informations for Master Browsers
workgroups = []
for masterbrowser in masters:
    for workgroup in lookuphost(masterbrowser,master=True):
        workgroups.append(workgroup)
print 'Avaliable workgroups: %s' % workgroups

# Lookup Hosts of Workgroups
workgroupshosts = {}
for workgroup in workgroups:
    def lookupworkgrouphosts(workgroup):
        """
        Lookup all hosts of workgroup
        """
        command = copy(workgroupcommand)
        command.append(workgroup)
        process = Popen(command, stdout=PIPE, stderr=PIPE)
        process.wait()
        stdout = process.stdout.read()
        process.stdout.close()
        match = masterRe.search(stdout)
        while match:
            if not workgroupshosts.has_key(workgroup):
                workgroupshosts[workgroup] = []
            workgroupshosts[workgroup].append(match.group(1))
            match = masterRe.search(stdout, match.end())
    lookupworkgrouphosts(workgroup)
print 'Workgroup / Host map: %s' % workgroupshosts

# Lookup Host Informations
workgroupshostsinfo = {}
for workgroup in workgroupshosts:
    if not workgroupshostsinfo.has_key(workgroup):
        workgroupshostsinfo[workgroup] = []
    for host in workgroupshosts[workgroup]:
        for hostinfo in lookuphost(host):
            hostdnsname =  gethostbyaddr(host)[0]
            hostdnsip = gethostbyname(hostinfo)
            workgroupshostsinfo[workgroup].append((host, hostdnsip, hostinfo, hostdnsname))
print 'Workgroup / IP / Hostname map: %s' % workgroupshostsinfo
