"""
Shodan CLI

Note: Always run "shodan init <api key>" before trying to execute any other command!

A simple interface to search Shodan, download data and parse compressed JSON files.
The following commands are currently supported:

    alert
    convert
    count
    data
    download
    honeyscore
    host
    info
    init
    myip
    parse
    radar
    scan
    search
    stats
    stream

"""

import click
import csv
import os
import os.path
import shodan
import shodan.helpers as helpers
import threading
import requests
import time

# The file converters that are used to go from .json.gz to various other formats
from shodan.cli.converter import CsvConverter, KmlConverter, GeoJsonConverter, ExcelConverter, ImagesConverter

# Constants
from shodan.cli.settings import SHODAN_CONFIG_DIR, COLORIZE_FIELDS

# Helper methods
from shodan.cli.helpers import async_spinner, get_api_key, escape_data, timestr, open_streaming_file, get_banner_field, match_filters
from shodan.cli.host import HOST_PRINT

# Allow 3rd-parties to develop custom commands
from click_plugins import with_plugins
from pkg_resources import iter_entry_points

# Make "-h" work like "--help"
CONTEXT_SETTINGS = dict(help_option_names=['-h', '--help'])

# Define a basestring type if necessary for Python3 compatibility
try:
    basestring
except NameError:
    basestring = str


# Define the main entry point for all of our commands
# and expose a way for 3rd-party plugins to tie into the Shodan CLI.
@with_plugins(iter_entry_points('shodan.cli.plugins'))
@click.group(context_settings=CONTEXT_SETTINGS)
def main():
    pass


# Large subcommands are stored in separate modules
from shodan.cli.alert import alert
from shodan.cli.data import data
from shodan.cli.organization import org
from shodan.cli.scan import scan
main.add_command(alert)
main.add_command(data)
main.add_command(org)
main.add_command(scan)


CONVERTERS = {
    'kml': KmlConverter,
    'csv': CsvConverter,
    'geo.json': GeoJsonConverter,
    'images': ImagesConverter,
    'xlsx': ExcelConverter,
}
@main.command()
@click.argument('input', metavar='<input file>')
@click.argument('format', metavar='<output format>', type=click.Choice(CONVERTERS.keys()))
def convert(input, format):
    """Convert the given input data file into a different format. The following file formats are supported:

    kml, csv, geo.json, images, xlsx

    Example: shodan convert data.json.gz kml
    """
    # Get the basename for the input file
    basename = input.replace('.json.gz', '').replace('.json', '')

    # Add the new file extension based on the format
    filename = '{}.{}'.format(basename, format)

    # Open the output file
    fout = open(filename, 'w')

    # Start a spinner
    finished_event = threading.Event()
    progress_bar_thread = threading.Thread(target=async_spinner, args=(finished_event,))
    progress_bar_thread.start()

    # Initialize the file converter
    converter = CONVERTERS.get(format)(fout)

    converter.process([input])

    finished_event.set()
    progress_bar_thread.join()

    if format == 'images':
        click.echo(click.style('\rSuccessfully extracted images to directory: {}'.format(converter.dirname), fg='green'))
    else:
        click.echo(click.style('\rSuccessfully created new file: {}'.format(filename), fg='green'))


@main.command()
@click.argument('key', metavar='<api key>')
def init(key):
    """Initialize the Shodan command-line"""
    # Create the directory if necessary
    shodan_dir = os.path.expanduser(SHODAN_CONFIG_DIR)
    if not os.path.isdir(shodan_dir):
        try:
            os.mkdir(shodan_dir)
        except OSError:
            raise click.ClickException('Unable to create directory to store the Shodan API key ({})'.format(shodan_dir))

    # Make sure it's a valid API key
    key = key.strip()
    try:
        api = shodan.Shodan(key)
        api.info()
    except shodan.APIError as e:
        raise click.ClickException(e.value)

    # Store the API key in the user's directory
    keyfile = shodan_dir + '/api_key'
    with open(keyfile, 'w') as fout:
        fout.write(key.strip())
        click.echo(click.style('Successfully initialized', fg='green'))

    os.chmod(keyfile, 0o600)

@main.command()
@click.argument('query', metavar='<search query>', nargs=-1)
def count(query):
    """Returns the number of results for a search"""
    key = get_api_key()

    # Create the query string out of the provided tuple
    query = ' '.join(query).strip()

    # Make sure the user didn't supply an empty string
    if query == '':
        raise click.ClickException('Empty search query')

    # Perform the search
    api = shodan.Shodan(key)
    try:
        results = api.count(query)
    except shodan.APIError as e:
        raise click.ClickException(e.value)

    click.echo(results['total'])


@main.command()
@click.option('--limit', help='The number of results you want to download. -1 to download all the data possible.', default=1000, type=int)
@click.argument('filename', metavar='<filename>')
@click.argument('query', metavar='<search query>', nargs=-1)
def download(limit, filename, query):
    """Download search results and save them in a compressed JSON file."""
    key = get_api_key()

    # Create the query string out of the provided tuple
    query = ' '.join(query).strip()

    # Make sure the user didn't supply an empty string
    if query == '':
        raise click.ClickException('Empty search query')

    filename = filename.strip()
    if filename == '':
        raise click.ClickException('Empty filename')

    # Add the appropriate extension if it's not there atm
    if not filename.endswith('.json.gz'):
        filename += '.json.gz'

    # Perform the search
    api = shodan.Shodan(key)

    try:
        total = api.count(query)['total']
        info = api.info()
    except:
        raise click.ClickException('The Shodan API is unresponsive at the moment, please try again later.')

    # Print some summary information about the download request
    click.echo('Search query:\t\t\t%s' % query)
    click.echo('Total number of results:\t%s' % total)
    click.echo('Query credits left:\t\t%s' % info['unlocked_left'])
    click.echo('Output file:\t\t\t%s' % filename)

    if limit > total:
        limit = total

    # A limit of -1 means that we should download all the data
    if limit <= 0:
        limit = total

    with helpers.open_file(filename, 'w') as fout:
        count = 0
        try:
            cursor = api.search_cursor(query, minify=False)
            with click.progressbar(cursor, length=limit) as bar:
                for banner in bar:
                    helpers.write_banner(fout, banner)
                    count += 1

                    if count >= limit:
                        break
        except Exception:
            pass

        # Let the user know we're done
        if count < limit:
            click.echo(click.style('Notice: fewer results were saved than requested', 'yellow'))
        click.echo(click.style(u'Saved {} results into file {}'.format(count, filename), 'green'))


@main.command()
@click.option('--format', help='The output format for the host information. Possible values are: pretty, tsv.', default='pretty', type=click.Choice(['pretty', 'tsv']))
@click.option('--history', help='Show the complete history of the host.', default=False, is_flag=True)
@click.option('--filename', '-O', help='Save the host information in the given file (append if file exists).', default=None)
@click.option('--save', '-S', help='Save the host information in the a file named after the IP (append if file exists).', default=False, is_flag=True)
@click.argument('ip', metavar='<ip address>')
def host(format, history, filename, save, ip):
    """View all available information for an IP address"""
    key = get_api_key()
    api = shodan.Shodan(key)

    try:
        host = api.host(ip, history=history)

        # Print the host information to the terminal using the user-specified format
        HOST_PRINT[format](host, history=history)

        # Store the results
        if filename or save:
            if save:
                filename = '{}.json.gz'.format(ip)

            # Add the appropriate extension if it's not there atm
            if not filename.endswith('.json.gz'):
                filename += '.json.gz'

            # Create/ append to the file
            fout = helpers.open_file(filename)

            for banner in sorted(host['data'], key=lambda k: k['port']):
                if 'placeholder' not in banner:
                    helpers.write_banner(fout, banner)
    except shodan.APIError as e:
        raise click.ClickException(e.value)



@main.command()
def info():
    """Shows general information about your account"""
    key = get_api_key()
    api = shodan.Shodan(key)
    try:
        results = api.info()
    except shodan.APIError as e:
        raise click.ClickException(e.value)

    click.echo("""Query credits available: {0}
Scan credits available: {1}
    """.format(results['query_credits'], results['scan_credits']))


@main.command()
@click.option('--color/--no-color', default=True)
@click.option('--fields', help='List of properties to output.', default='ip_str,port,hostnames,data')
@click.option('--filters', '-f', help='Filter the results for specific values using key:value pairs.', multiple=True)
@click.option('--filename', '-O', help='Save the filtered results in the given file (append if file exists).')
@click.option('--separator', help='The separator between the properties of the search results.', default=u'\t')
@click.argument('filenames', metavar='<filenames>', type=click.Path(exists=True), nargs=-1)
def parse(color, fields, filters, filename, separator, filenames):
    """Extract information out of compressed JSON files."""
    # Strip out any whitespace in the fields and turn them into an array
    fields = [item.strip() for item in fields.split(',')]

    if len(fields) == 0:
        raise click.ClickException('Please define at least one property to show')

    has_filters = len(filters) > 0


    # Setup the output file handle
    fout = None
    if filename:
        # If no filters were provided raise an error since it doesn't make much sense w/out them
        if not has_filters:
            raise click.ClickException('Output file specified without any filters. Need to use filters with this option.')

        # Add the appropriate extension if it's not there atm
        if not filename.endswith('.json.gz'):
            filename += '.json.gz'
        fout = helpers.open_file(filename)

    for banner in helpers.iterate_files(filenames):
        row = u''

        # Validate the banner against any provided filters
        if has_filters and not match_filters(banner, filters):
            continue

        # Append the data
        if fout:
            helpers.write_banner(fout, banner)

        # Loop over all the fields and print the banner as a row
        for field in fields:
            tmp = u''
            value = get_banner_field(banner, field)
            if value:
                field_type = type(value)

                # If the field is an array then merge it together
                if field_type == list:
                    tmp = u';'.join(value)
                elif field_type in [int, float]:
                    tmp = u'{}'.format(value)
                else:
                    tmp = escape_data(value)

                # Colorize certain fields if the user wants it
                if color:
                    tmp = click.style(tmp, fg=COLORIZE_FIELDS.get(field, 'white'))

                # Add the field information to the row
                row += tmp
            row += separator

        click.echo(row)


@main.command()
def myip():
    """Print your external IP address"""
    key = get_api_key()

    api = shodan.Shodan(key)
    try:
        click.echo(api.tools.myip())
    except shodan.APIError as e:
        raise click.ClickException(e.value)


@main.command()
@click.option('--color/--no-color', default=True)
@click.option('--fields', help='List of properties to show in the search results.', default='ip_str,port,hostnames,data')
@click.option('--limit', help='The number of search results that should be returned. Maximum: 1000', default=100, type=int)
@click.option('--separator', help='The separator between the properties of the search results.', default='\t')
@click.argument('query', metavar='<search query>', nargs=-1)
def search(color, fields, limit, separator, query):
    """Search the Shodan database"""
    key = get_api_key()

    # Create the query string out of the provided tuple
    query = ' '.join(query).strip()

    # Make sure the user didn't supply an empty string
    if query == '':
        raise click.ClickException('Empty search query')

    # For now we only allow up to 1000 results at a time
    if limit > 1000:
        raise click.ClickException('Too many results requested, maximum is 1,000')

    # Strip out any whitespace in the fields and turn them into an array
    fields = [item.strip() for item in fields.split(',')]

    if len(fields) == 0:
        raise click.ClickException('Please define at least one property to show')

    # Perform the search
    api = shodan.Shodan(key)
    try:
        results = api.search(query, limit=limit)
    except shodan.APIError as e:
        raise click.ClickException(e.value)

    # Error out if no results were found
    if results['total'] == 0:
        raise click.ClickException('No search results found')

    # We buffer the entire output so we can use click's pager functionality
    output = u''
    for banner in results['matches']:
        row = u''

        # Loop over all the fields and print the banner as a row
        for field in fields:
            tmp = u''
            value = get_banner_field(banner, field)
            if value:
                field_type = type(value)

                # If the field is an array then merge it together
                if field_type == list:
                    tmp = u';'.join(value)
                elif field_type in [int, float]:
                    tmp = u'{}'.format(value)
                else:
                    tmp = escape_data(value)

                # Colorize certain fields if the user wants it
                if color:
                    tmp = click.style(tmp, fg=COLORIZE_FIELDS.get(field, 'white'))

                # Add the field information to the row
                row += tmp
            row += separator

            # click.echo(out + separator, nl=False)
        output += row + u'\n'
        # click.echo('')
    click.echo_via_pager(output)


@main.command()
@click.option('--limit', help='The number of results to return.', default=10, type=int)
@click.option('--facets', help='List of facets to get statistics for.', default='country,org')
@click.option('--filename', '-O', help='Save the results in a CSV file of the provided name.', default=None)
@click.argument('query', metavar='<search query>', nargs=-1)
def stats(limit, facets, filename, query):
    """Provide summary information about a search query"""
    # Setup Shodan
    key = get_api_key()
    api = shodan.Shodan(key)

    # Create the query string out of the provided tuple
    query = ' '.join(query).strip()

    # Make sure the user didn't supply an empty string
    if query == '':
        raise click.ClickException('Empty search query')

    facets = facets.split(',')
    facets = [(facet, limit) for facet in facets]

    # Perform the search
    try:
        results = api.count(query, facets=facets)
    except shodan.APIError as e:
        raise click.ClickException(e.value)

    # Print the stats tables
    for facet in results['facets']:
        click.echo('Top {} Results for Facet: {}'.format(len(results['facets'][facet]), facet))

        for item in results['facets'][facet]:
            # Force the value to be a string - necessary because some facet values are numbers
            value = u'{}'.format(item['value'])

            click.echo(click.style(u'{:28s}'.format(value), fg='cyan'), nl=False)
            click.echo(click.style(u'{:12,d}'.format(item['count']), fg='green'))

        click.echo('')

    # Create the output file if requested
    fout = None
    if filename:
        if not filename.endswith('.csv'):
            filename += '.csv'
        fout = open(filename, 'w')
        writer = csv.writer(fout, dialect=csv.excel)

        # Write the header
        writer.writerow(['Query', query])

        # Add an empty line to separate rows
        writer.writerow([])

        # Write the header that contains the facets
        row = []
        for facet in results['facets']:
            row.append(facet)
            row.append('')
        writer.writerow(row)

        # Every facet has 2 columns (key, value)
        counter = 0
        has_items = True
        while has_items:
            # pylint: disable=W0612
            row = ['' for i in range(len(results['facets']) * 2)]

            pos = 0
            has_items = False
            for facet in results['facets']:
                values = results['facets'][facet]

                # Add the values for the facet into the current row
                if len(values) > counter:
                    has_items = True
                    row[pos] = values[counter]['value']
                    row[pos+1] = values[counter]['count']

                pos += 2

            # Write out the row
            if has_items:
                writer.writerow(row)

            # Move to the next row of values
            counter += 1


@main.command()
@click.option('--color/--no-color', default=True)
@click.option('--fields', help='List of properties to output.', default='ip_str,port,hostnames,data')
@click.option('--separator', help='The separator between the properties of the search results.', default='\t')
@click.option('--limit', help='The number of results you want to download. -1 to download all the data possible.', default=-1, type=int)
@click.option('--datadir', help='Save the stream data into the specified directory as .json.gz files.', default=None, type=str)
@click.option('--ports', help='A comma-separated list of ports to grab data on.', default=None, type=str)
@click.option('--quiet', help='Disable the printing of information to the screen.', is_flag=True)
@click.option('--timeout', help='Timeout. Should the shodan stream cease to send data, then timeout after <timeout> seconds.', default=0, type=int)
@click.option('--streamer', help='Specify a custom Shodan stream server to use for grabbing data.', default='https://stream.shodan.io', type=str)
@click.option('--countries', help='A comma-separated list of countries to grab data on.', default=None, type=str)
@click.option('--asn', help='A comma-separated list of ASNs to grab data on.', default=None, type=str)
@click.option('--alert', help='The network alert ID or "all" to subscribe to all network alerts on your account.', default=None, type=str)
@click.option('--compresslevel', help='The gzip compression level (0-9; 0 = no compression, 9 = most compression', default=9, type=int)
def stream(color, fields, separator, limit, datadir, ports, quiet, timeout, streamer, countries,  asn, alert, compresslevel):
    """Stream data in real-time."""
    # Setup the Shodan API
    key = get_api_key()
    api = shodan.Shodan(key)

    # Temporarily change the baseurl
    api.stream.base_url = streamer

    # Strip out any whitespace in the fields and turn them into an array
    fields = [item.strip() for item in fields.split(',')]

    if len(fields) == 0:
        raise click.ClickException('Please define at least one property to show')

    # The user must choose "ports", "countries", "asn" or nothing - can't select multiple
    # filtered streams at once.
    stream_type = []
    if ports:
        stream_type.append('ports')
    if countries:
        stream_type.append('countries')
    if asn:
        stream_type.append('asn')
    if alert:
        stream_type.append('alert')

    if len(stream_type) > 1:
        raise click.ClickException('Please use --ports, --countries OR --asn. You cant subscribe to multiple filtered streams at once.')

    stream_args = None

    # Turn the list of ports into integers
    if ports:
        try:
            stream_args = [int(item.strip()) for item in ports.split(',')]
        except ValueError:
            raise click.ClickException('Invalid list of ports')

    if alert:
        alert = alert.strip()
        if alert.lower() != 'all':
            stream_args = alert

    if asn:
        stream_args = asn.split(',')

    if countries:
        stream_args = countries.split(',')

    # Flatten the list of stream types
    # Possible values are:
    # - all
    # - asn
    # - countries
    # - ports
    if len(stream_type) == 1:
        stream_type = stream_type[0]
    else:
        stream_type = 'all'

    # Decide which stream to subscribe to based on whether or not ports were selected
    def _create_stream(name, args, timeout):
        return {
            'all': api.stream.banners(timeout=timeout),
            'alert': api.stream.alert(args, timeout=timeout),
            'asn': api.stream.asn(args, timeout=timeout),
            'countries': api.stream.countries(args, timeout=timeout),
            'ports': api.stream.ports(args, timeout=timeout),
        }.get(name, 'all')

    stream = _create_stream(stream_type, stream_args, timeout=timeout)

    counter = 0
    quit = False
    last_time = timestr()
    fout = None

    if datadir:
        fout = open_streaming_file(datadir, last_time, compresslevel)

    while not quit:
        try:
            for banner in stream:
                # Limit the number of results to output
                if limit > 0:
                    counter += 1

                    if counter > limit:
                        quit = True
                        break

                # Write the data to the file
                if datadir:
                    cur_time = timestr()
                    if cur_time != last_time:
                            last_time = cur_time
                            fout.close()
                            fout = open_streaming_file(datadir, last_time)
                    helpers.write_banner(fout, banner)

                # Print the banner information to stdout
                if not quiet:
                    row = u''

                    # Loop over all the fields and print the banner as a row
                    for field in fields:
                        tmp = u''
                        value = get_banner_field(banner, field)
                        if value:
                            field_type = type(value)

                            # If the field is an array then merge it together
                            if field_type == list:
                                tmp = u';'.join(value)
                            elif field_type in [int, float]:
                                tmp = u'{}'.format(value)
                            else:
                                tmp = escape_data(value)

                            # Colorize certain fields if the user wants it
                            if color:
                                tmp = click.style(tmp, fg=COLORIZE_FIELDS.get(field, 'white'))

                            # Add the field information to the row
                            row += tmp
                        row += separator

                    click.echo(row)
        except requests.exceptions.Timeout:
            raise click.ClickException('Connection timed out')
        except KeyboardInterrupt:
            quit = True
        except shodan.APIError as e:
            raise click.ClickException(e.value)
        except Exception:
            # For other errors lets just wait a bit and try to reconnect again
            time.sleep(1)

            # Create a new stream object to subscribe to
            stream = _create_stream(stream_type, stream_args, timeout=timeout)


@main.command()
@click.argument('ip', metavar='<IP address>')
def honeyscore(ip):
    """Check whether the IP is a honeypot or not."""
    key = get_api_key()
    api = shodan.Shodan(key)

    try:
        score = api.labs.honeyscore(ip)

        if score == 1.0:
            click.echo(click.style('Honeypot detected', fg='red'))
        elif score > 0.5:
            click.echo(click.style('Probably a honeypot', fg='yellow'))
        else:
            click.echo(click.style('Not a honeypot', fg='green'))

        click.echo('Score: {}'.format(score))
    except:
        raise click.ClickException('Unable to calculate honeyscore')


@main.command()
def radar():
    """Real-Time Map of some results as Shodan finds them."""
    key = get_api_key()
    api = shodan.Shodan(key)

    from shodan.cli.worldmap import launch_map

    try:
        launch_map(api)
    except shodan.APIError as e:
        raise click.ClickException(e.value)
    except Exception as e:
        raise click.ClickException(u'{}'.format(e))

if __name__ == '__main__':
    main()
