#!/usr/bin/env python
# By Steve 'Ashcrow' Milner

import nmap

try:
    from multiprocessing import Process, Queue
except ImportError:
    # For pre 2.6 releases
    from threading import Thread as Process
    from Queue import Queue


class NmapProcess(object):
    """
    Process scans in a non-blocking way.
    """

    def __init__(self, host):
        """
        Creates an instance.

        :Parameters:
           - `host`: host to scan
        """
        self.__host = host
        self.__queue = Queue()
        self.__nm = nmap.PortScanner()
        self.__process = None

    def scan(self, ports=None, arguments='-sV', callback=None):
        """
        Starts the scan taking the same input as nmap scan except the host.

        :Parameters:
           - `ports`: ports to scan
           - `arguments` arguments for nmap
           - `callback`: optional callback to execute when results are available
        """

        def scan_process(ports, arguments, callback=None):
            self.__nm.scan(self.__host, ports, arguments)
            self.__queue.put(self.__nm)
            if callback and callable(callback):
                callback(self.__nm)

        self.__process = Process(
            target=scan_process, args=(ports, arguments, callback))
        self.__process.daemon = True
        self.__process.start()

    def __get_results(self):
        """
        Returns results if they are available, else None.
        """
        if self.__queue.empty():
            return None
        return self.__queue.get()[self.__host]

    def __del__(self):
        """
        Clean up.
        """
        self.__process.join()

    # Properties
    results = property(__get_results)
    results_available = property(lambda s: not s.__process.is_alive())


if __name__ == '__main__':
    n = NmapProcess('127.0.0.1')

    def p(s):
        print("From Callback: " + str(s['127.0.0.1']))
    n.scan(arguments="-sV -T4 -A", callback=p)
    # Do stuff here ...
    from time import sleep
    while not n.results_available:
        print("Waiting ...")
        sleep(3)
    print("From property: " + str(n.results))

description = [[
Checks systems for Arugizer, the energizer bunny trojan.
]]

---
-- @usage
-- nmap --script arugizer.nse <target>
-- @output
-- Host script results:
-- |_ arugizer: Infected


author = "Steve Milner"
license = "Same as Nmap--See http://nmap.org/book/man-legal.html"
categories = {"default", "safe", "discovery"}

require("nmap")

-- File that holds the following as found by ron@skullsecurity.net
-- Hex: \xC2\xE5\xE5\xE5\x9E\xA0\xD7\xA4\xA6\xD0\xD5\xDD\xDC\xC8\xD6\xDD\xD7\xD5\xC8\xD1\xD6\x83\x80\xC8\xDD\xA4\xD1\xA1\xC8\xA4\xD2\xD5\xD7\xDD\xA3\xA4\xA1\xDD\xA6\xD7\xDD\x98\xE5
-- ASCII: {E2AC5089-3820-43fe-8A4D-A7028FAD8C28}
f = assert(io.open(nmap.fetchfile('nselib/data/arugizer_ping.data'), "rb"))
ping = f:read("*all")
f:close()

-- File that holds the following as found by ron@skullsecurity.net
-- Hex: \xbc\xa0\xb6
-- ASCII: YES
f = assert(io.open(nmap.fetchfile('nselib/data/arugizer_yes.data'), "rb"))
yes = f:read("*all")
f:close()


--- Rule to decide if the action should take place.
-- We only trigger if the port is 7777
-- @param host The host table from nmap.
-- @param port The port info from nmap.
portrule = function(host, port)
    if port.state == "open" and
       port.protocol == "tcp" and
       port.number == 7777 then
        return true
    end
    return false
end


--- Takes action if the host rule is triggered.
-- Checks if port 7777 responds to known string
-- @param host The host table from nmap.
-- @param port The port info from nmap.
action = function(host, port)
    local socket = nmap.new_socket()
    socket:set_timeout(1000)
    socket:connect(host.ip, 7777)

    socket:send(ping)
    local status, response = socket:receive()
    socket:close()

    if (string.sub(response, 0, 3) == yes) then
        return "Infected"
    end
    return "Seems OK ..."
end