#! /usr/bin/env python
import json
from dateutil.parser import parse
import shodan
from harpoon.commands.base import Command, Subcommand
from harpoon.lib.utils import unbracket


class SubcommandIP(Subcommand):
    description = "Get information on an IP address"
    cmd = "ip"

    def add_arguments(self, parser):
        parser.add_argument('IP', help='IP to be searched')
        parser.add_argument(
            '--history', '-H', action='store_true',
            help='Also display historical information')
        parser.add_argument(
            '-v', '--verbose', action='store_true',
            help="Verbose mode (display raw json)")
        parser.add_argument(
            '-s', '--summary', action='store_true',
            help="Only display information for ports 22, 80 and 443")

    def run(self, args):
        api = shodan.Shodan(self._config_data['Shodan']['key'])
        try:
            self.data = api.host(args.IP, history=args.history)
        except shodan.exception.APIError:
            self.data = None

    def display(self, args):
        if self.data is None:
            print("IP not found in Shodan")
            return

        if args.verbose:
            print(json.dumps(self.data, indent=4))
        else:
            if args.summary:
                for d in self.data['data']:
                    if d['port'] == 22:
                        print("{} - port 22 ssh - {}".format(
                                d['timestamp'][:19],
                                d['data'].split("\n")[0]
                            )
                        )
                    elif d['port'] == 80:
                        print("{} - port 80 http - Server \"{}\"" .format(
                                d['timestamp'][:19],
                                d['http']['server']
                            )
                        )
                    elif d['port'] == 443:
                        if 'cert' in d['ssl']:
                            print("%s - port 443 https - Cert \"%s\" \"%s\" %s - Server \"%s\"" % (
                                    d['timestamp'][:19],
                                    d['ssl']['cert']['subject']['CN'],
                                    d['ssl']['cert']['issuer']['CN'],
                                    d['ssl']['cert']['fingerprint']['sha1'],
                                    d['http']['server']
                                )
                            )
                        else:
                            print("%s - port 443 https - Cert Unknown- Server \"%s\"" % (
                                    d['timestamp'][:19],
                                    d['http']['server']
                                )
                            )
            else:
                for d in self.data['data']:
                    print(d['timestamp'])
                    print(d['_shodan']['module'])
                    print("%s/%i" % (d['transport'], d['port']))
                    print(d['data'])
                    if 'html' in d:
                        print(d['html'][:2000])
                    if 'http' in d:
                        print(json.dumps(d['http'])[:3000])
                    print('')


class SubcommandSearch(Subcommand):
    description = "Search in shodan"
    cmd = "search"

    def add_arguments(self, parser):
        parser.add_argument('QUERY', help='Query')
        parser.add_argument(
            '-v', '--verbose', action='store_true',
            help="Verbose mode (display raw json)")

    def run(self, args):
        api = shodan.Shodan(self._config_data['Shodan']['key'])
        self.data = api.search(args.QUERY)

    def display(self, args):
        if args.verbose:
            print(json.dumps(self.data, indent=4))
        else:
            print('%i results' % self.data['total'])
            for r in self.data['matches']:
                print('[+] {} ({}): port {}/{} -> {}\n'.format(
                        r['ip_str'],
                        r['org'],
                        r['transport'],
                        r['port'],
                        r['data'][:1000]
                    )
                )


class SubcommandSsh(Subcommand):
    description = "Get history of ssh keys from Shodan historical data"
    cmd = "ssh"

    def add_arguments(self, parser):
        parser.add_argument('IP', help='IP address')

    def run(self, args):
        api = shodan.Shodan(self._config_data['Shodan']['key'])
        try:
            res = api.host(unbracket(args.IP), history=True)
        except shodan.exception.APIError as e:
            print(e)
            self.data = None
            return

        self.data = {}
        for event in res['data']:
            if event['_shodan']['module'] == 'ssh':
                if 'ssh' not in event:
                    continue
                fingerprint = event['ssh']['fingerprint']
                date = parse(event['timestamp'])
                if fingerprint not in self.data:
                    self.data[fingerprint] = {
                        'first': date,
                        'last': date,
                        'fingerprint': fingerprint
                    }
                else:
                    if self.data[fingerprint]['first'] > date:
                        self.data[fingerprint]['first'] = date
                    if self.data[fingerprint]['last'] < date:
                        self.data[fingerprint]['last'] = date

    def display(self, args):
        if self.data is None:
            print("IP not found in Shodan")
            return

        for val in sorted(self.data.values(), key=lambda x: x['first']):
            print('{} - {} -> {}'.format(
                val['fingerprint'],
                val['first'].strftime('%Y-%m-%d'),
                val['last'].strftime('%Y-%m-%d')
            ))


class SubcommandQuota(Subcommand):
    description = "Show Shodan account quota information"
    cmd = "quota"

    def run(self, args):
        api = shodan.Shodan(self._config_data['Shodan']['key'])
        self.data = api.info()

    def display(self, args):
        print(json.dumps(self.data, indent=4))


# ------------------------------------ Main Command ---------------------------
class CommandShodan(Command):
    """
    # Shodan

    **Queries information from shodan.io API***

    * Get information on an IP : `harpoon shodan ip IP`
    * Get summary (only ports 22, 80 and 443) of historical data on an ip : `harpoon shodan ip -H -s IP`
    * Get raw json of historical data : `harpoon shodan ip -H -v IP`
    * Search in the database: `harpoon shodan search SEARCH`
    """
    name = "shodan"
    description = "Requests Shodan API"
    config = {'Shodan': ['key']}

    def __init__(self, config):
        super().__init__(config=config)
        self.add_subcommand(SubcommandIP)
        self.add_subcommand(SubcommandSearch)
        self.add_subcommand(SubcommandSsh)
        self.add_subcommand(SubcommandQuota)

    def intel_ip(self, query, data, conf):
        print("[+] Checking Shodan...")
        api = shodan.Shodan(self._config_data['Shodan']['key'])
        try:
            res = api.host(query)
        except shodan.exception.APIError:
            pass
        else:
            for p in res["ports"]:
                data["ports"].append({
                    "port": p,
                    "source": "Shodan",
                    "info": ""
                })
