try:
    from json       import dumps, loads
except:
    from simplejson import dumps, loads
from urllib2    import urlopen
from urllib     import urlencode

__all__ = ['WebAPI']

class WebAPIError(Exception):
    def __init__(self, value):
        self.value = value
    
    def __str__(self):
        return self.value


class WebAPI:
    """Wrapper around the SHODAN webservices API"""
    
    class DatalossDb:
        
        def __init__(self, parent):
            self.parent = parent
        
        def search(self, **kwargs):
            """Search the Dataloss DB archive.
    
            Arguments:
            name          -- Name of the affected company/ organisation
            
            arrest        -- whether the incident resulted in an arrest
            breaches      -- the type of breach that occurred (Hack, MissingLaptop etc.)
            country       -- country where the incident took place
            ext           -- whether an external, third party was affected
            ext_names     -- the name of the third party company that was affected
            lawsuit       -- whether the incident resulted in a lawsuit
            records       -- the number of records that were lost/ stolen
            recovered     -- whether the affected items were recovered
            sub_types     -- the sub-categorization of the affected company/ organization
            source        -- whether the incident occurred from inside or outside the organization
            stocks        -- stock symbol of the affected company
            types         -- the basic type of organization (government, business, educational)
            uid           -- unique ID for the incident
    
            Returns:
            A dictionary with 2 main items: matches (list) and total (int).
    
            """
            return self.parent._request('datalossdb/search', dict(**kwargs))
    
    class Exploits:
        
        def __init__(self, parent):
            self.parent = parent
            
        def search(self, query, sources=[], cve=None, osvdb=None, msb=None, bid=None):
            """Search the entire Shodan Exploits archive using the same query syntax
            as the website.
            
            Arguments:
            query    -- exploit search query; same syntax as website
            
            Optional arguments:
            sources  -- metasploit, cve, osvdb, exploitdb, or packetstorm
            cve      -- CVE identifier (ex. 2010-0432)
            osvdb    -- OSVDB identifier (ex. 11666)
            msb      -- Microsoft Security Bulletin ID (ex. MS05-030)
            bid      -- Bugtraq identifier (ex. 13951)
            
            """
            if sources:
                query += ' source:' + ','.join(sources)
            if cve:
                query += ' cve:%s' % (str(cve).strip())
            if osvdb:
                query += ' osvdb:%s' % (str(osvdb).strip())
            if msb:
                query += ' msb:%s' % (str(msb).strip())
            if bid:
                query += ' bid:%s' % (str(bid).strip())
            return self.parent._request('search_exploits', {'q': query})
    
    class ExploitDb:
        
        def __init__(self, parent):
            self.parent = parent
        
        def download(self, id):
            """Download the exploit code from the ExploitDB archive.
    
            Arguments:
            id    -- ID of the ExploitDB entry
    
            Returns:
            A dictionary with the following fields:
            filename        -- Name of the file
            content-type    -- Mimetype
            data            -- Contents of the file
    
            """
            return self.parent._request('exploitdb/download', {'id': id})
        
        def search(self, query, **kwargs):
            """Search the ExploitDB archive.
    
            Arguments:
            query     -- Search terms
            
            Optional arguments:
            author    -- Name of the exploit submitter
            platform  -- Target platform (e.g. windows, linux, hardware etc.)
            port      -- Service port number
            type      -- Any, dos, local, papers, remote, shellcode and webapps
    
            Returns:
            A dictionary with 2 main items: matches (list) and total (int).
            Each item in 'matches' is a dictionary with the following elements:
            
            id
            author
            date
            description
            platform
            port
            type
    
            """
            return self.parent._request('exploitdb/search', dict(q=query, **kwargs))
    
    class Msf:
        
        def __init__(self, parent):
            self.parent = parent
            
        def download(self, id):
            """Download a metasploit module given the fullname (id) of it.
            
            Arguments:
            id        -- fullname of the module (ex. auxiliary/admin/backupexec/dump)
            
            Returns:
            A dictionary with the following fields:
            filename        -- Name of the file
            content-type    -- Mimetype
            data            -- File content
            """
            return self.parent._request('msf/download', {'id': id})
        
        def search(self, query, **kwargs):
            """Search for a Metasploit module.
            """
            return self.parent._request('msf/search', dict(q=query, **kwargs))
    
    def __init__(self, key):
        """Initializes the API object.
        
        Arguments:
        key -- your API key
        
        """
        self.api_key = key
        self.base_url = 'http://www.shodanhq.com/api/'
        self.dataloss = self.DatalossDb(self)
        self.exploits = self.Exploits(self)
        self.exploitdb = self.ExploitDb(self)
        self.msf = self.Msf(self)
    
    def _request(self, function, params):
        """General-purpose function to create web requests to SHODAN.
        
        Arguments:
        function  -- name of the function you want to execute
        params      -- dictionary of parameters for the function
        
        Returns
        A JSON string containing the function's results.
        
        """
        # Add the API key parameter automatically
        params['key'] = self.api_key
        
        # Send the request
        data = urlopen(self.base_url + function + '?' + urlencode(params)).read()
        
        # Parse the text into JSON
        data = loads(data)
        
        # Raise an exception if an error occurred
        if data.get('error', None):
            raise WebAPIError(data['error'])
        
        # Return the data
        return data
    
    def fingerprint(self, banner):
        """Determine the software based on the banner.
        
        Arguments:
        banner  - HTTP banner
        
        Returns:
        A list of software that matched the given banner.
        """
        return self._request('fingerprint', {'banner': banner})
    
    def host(self, ip):
        """Get all available information on an IP.

        Arguments:
        ip    -- IP of the computer

        Returns:
        All available information SHODAN has on the given IP,
        subject to API key restrictions.

        """
        return self._request('host', {'ip': ip})
    
    def search(self, query):
        """Search the SHODAN database.
        
        Arguments:
        query    -- search query; identical syntax to the website
        
        Returns:
        A dictionary with 3 main items: matches, countries and total.
        Visit the website for more detailed information.
        
        """
        return self._request('search', {'q': query})
