# nmblookup.py
# -*- coding:utf-8 -*-
# Copyright (C) 2008-2009 Stefan J. Betz <stefan_betz@gmx.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/>.

"""
==============================
Network Message Block Handling
==============================

This part is usefull for basic network browsing tasks like:

* List avaliable workgroups
* List hosts in a workgroup
* Query the master browser of a workgroup
* Scan for NetBIOS Hostname and Workgroup
"""

from re import compile
from copy import copy
from subprocess import Popen, PIPE
from pyneighborhood import db, dblock
from socket import gethostbyaddr

masterRe = compile("\n((?:\d{1,3}\.){3}\d{1,3})")
lookupRe = compile("\n\t([^\s]*)\s*(<(?:20|1e)>)")
workgroupcommand = [ "nmblookup" ]
lookupcommand = [ "nmblookup", "-A" ]
mastercommand = [ "nmblookup", "-M", "--", "-" ]

def lookup(ip):
    """
    Lookups NetBIOS name and workgroup through nmblookup.
    
    :param ip: IPv4 of a specific host to query
    :returns: A tuple with name and workgroup
    :rtype: tuple
    """
    command = copy(lookupcommand)
    command.append(ip)
    process = Popen( command, executable = command[0],
                     stdout = PIPE,
                     stderr = PIPE )
    process.wait()
    output = process.stdout.read()
    match = lookupRe.search(output)
    
    # preinitialize name & workgroup variables
    # to avoid reference before assignment!
    name = None
    workgroup = None
    
    while match:
        if match.group(2) == "<20>":
            name = match.group(1)
        elif match.group(2) == "<1e>":
            workgroup = match.group(1)
        match = lookupRe.search(output, match.end())
    del match
    return (name, workgroup)

def query_masters():
    """
    A Generator which queries all Master Browsers.

    :returns: None
    :rtype: None
    """
    process = Popen(mastercommand, stdout=PIPE)
    process.wait()
    stdout = process.stdout.read()
    process.stdout.close()
    match = masterRe.search(stdout)
    while match:
        yield match.group(1)
        match = masterRe.search(stdout, match.end())

def query_workgroup(ip, manual=False):
    """
    Resolve workgroup of the host `ip` and add workgroup to database.

    :param ip: IPv4 of a specific host to query
    :param manual: ``True`` if this is an entry created by the user, ``False`` if this entry is created by scanning.
    :returns: Workgroup and Database-ID of this workgroup
    :rtype: tuple
    """
    # Step 1: nmblookup -A ip (hostlookupcommand)
    # Step 2: Check if workgroup already in database
    manual = 1 if manual else 0
    name, workgroup = lookup(ip)
    cursor = db.cursor()
    result = cursor.execute("""SELECT id FROM workgroups WHERE name=?;""", (workgroup,)).fetchall()
    if not result:
        dblock.acquire()
        cursor.execute("""INSERT INTO workgroups (name,manual) VALUES (?,?);""", (workgroup,manual,))
        dblock.release()
    elif len(result) == 1:
        dblock.acquire()
        cursor.execute("""UPDATE workgroups SET validated = 1 WHERE id = ?""", (result[0][0],))
        dblock.release()
    cursor.execute("""SELECT * FROM workgroups WHERE name=?;""", (workgroup,))
    return (workgroup, cursor.fetchone()[0])

def query_workgroup_hosts(workgroup):
    """
    Query all hosts of `workgroup`.

    :param workgroup: NetBIOS workgroup name
    :returns: Generator for all hosts in this workgroup.
    :rtype: generator
    """
    # Step 1: nmblookup workgroup (workgrouplookupcommand)
    command = copy(workgroupcommand)
    command.append(workgroup)
    process = Popen(command, stdout=PIPE)
    process.wait()
    stdout = process.stdout.read()
    process.stdout.close()
    match = masterRe.search(stdout)
    while match:
        yield match.group(1)
        match = masterRe.search(stdout, match.end())

def query_host(ip, manual=False):
    """
    Query the host `ip` and added it to the database.

    :param ip: IPv4 of host to query
    :param manual: ``True`` if this is an entry created by the user, ``False`` if this entry is created by scanning.
    :returns: None
    :rtype: None
    """
    # Step 1: Check if Host is already in Database
    manual = 1 if manual else 0
    cursor = db.cursor()
    result = cursor.execute("""SELECT id,ip,name FROM hosts WHERE ip = ?""", ( ip, )).fetchall()
    if len(result) == 1:
        dblock.acquire()
        cursor.execute("""UPDATE hosts SET validated = 1 WHERE id = ?""", (result[0][0],))
        dblock.release()
        return
    # Step 2: Create a new Database Entry...
    name, workgroup = lookup(ip)
    if not name or not workgroup:
        return
    workgroupid = cursor.execute("""SELECT id FROM workgroups WHERE name = ?""", ( workgroup, )).fetchone()[0]
    try:
        hostname = gethostbyaddr(ip)[0]
    except:
        hostname = ""
    dblock.acquire()
    cursor.execute("""INSERT INTO hosts ( ip,name,workgroup,hostname,manual ) VALUES ( ?, ?, ?, ?, ? )""", ( ip, name, workgroupid, hostname, manual, ))
    dblock.release()
