#!/usr/bin/python
import os
import sys
import logging
import click
import time
import simplejson as json
import base64
import requests
import threading 
import time
from threading import Thread, current_thread
import random
from random import randint
from pprint import pprint

"""
Documentation:

    This tool provides capability to run API scan 

    Example: apiclient.py scan  -k <key> -s <server> --url <domain/url>


REST API
--------------------------------------------------------------------------------
POST /api/v3/<api-key>/url/scan/<domain-name>[.json|.xml|.yaml]          Submits domain for new scan
POST /api/v3/<api-key>/url/scan/<base64-encoded-url>[.json|.xml|.yaml]   Submits domain for new scan
curl -X POST http://api.quttera.com/api/v3/A58E3D9E-514B-31DE-9617-2EB82B6550D0/url/scan/quttera.com.xml
--------------------------------------------------------------------------------
PUT /api/v3/<api-key>/url/scan/<domain-name>[.json|.xml|.yaml]           Resubmits domain, with existing threat report, for re-scan
PUT /api/v3/<api-key>/url/scan/<base64-encoded-url>[.json|.xml|.yaml]    Resubmits domain, with existing threat report, for re-scan
curl -X PUT -H 'Content-Length: 0' http://api.quttera.com/api/v3/A58E3D9E-514B-31DE-9617-2EB82B6550D0/url/scan/quttera.com.yaml
--------------------------------------------------------------------------------
GET /api/v3/<api-key>/url/status/<domain-name>[.json|.xml|.yaml]         Returns current status of domain submitted for scan
GET /api/v3/<api-key>/url/status/<base64-encoded-url>[.json|.xml|.yaml]  Returns current status of domain submitted for scan
curl http://api.quttera.com/api/v3/A58E3D9E-514B-31DE-9617-2EB82B6550D0/url/status/quttera.com.yaml
--------------------------------------------------------------------------------
GET /api/v3/<api-key>/url/report/<domain-name>[.json|.xml|.yaml]         Returns malware scan report for domain
GET /api/v3/<api-key>/url/report/<base64-encoded-url>[.json|.xml|.yaml]  Returns malware scan report for domain
curl http://api.quttera.com/api/v3/A58E3D9E-514B-31DE-9617-2EB82B6550D0/url/report/quttera.com.yaml
--------------------------------------------------------------------------------
GET /api/v3/<api-key>/blacklist/status/<domain-name>[.json|.xml|.yaml]         Returns current blacklist status of the domain
GET /api/v3/<api-key>/blacklist/status/<base64-encoded-url>[.json|.xml|.yaml]  Returns current blacklist status of the domain
curl http://api.quttera.com/api/v3/A58E3D9E-514B-31DE-9617-2EB82B6550D0/blacklist/status/quttera.com.yaml
--------------------------------------------------------------------------------
"""

logger = logging.getLogger()
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter(
    '[%(asctime)s:%(filename)s:%(lineno)s] %(levelname)s  %(message)s')

sth = logging.StreamHandler()
sth.setLevel(logging.DEBUG)
sth.setFormatter(formatter)
logger.addHandler(sth)
logfile = 'api_scale_test.log'

fileHandler = logging.FileHandler(logfile)
fileHandler.setFormatter(formatter)
logger.addHandler(fileHandler)


def log(severity, msg):
    logger.info("{} {} {}".format(current_thread().getName(),
                severity,
                msg))


def scan_request(url, key, server):
    request_url = "http://{}/api/v3/{}/url/scan/{}.json".format(
            server,
            key,
            url)

    print "Executing [POST {}]".format(request_url)
    try:
        response = requests.post(request_url)
        response.raise_for_status()
        return response.json()
    except:
        log("ERROR", "POST on {} failed".format(request_url))
        return None


def status_request(url, key, server):
    request_url = "http://{}/api/v3/{}/url/status/{}.json".format(
            server,
            key,
            url)

    print "Executing [GET {}]".format(request_url)
    try:
        response = requests.get(request_url)
        response.raise_for_status()
        return response.json()
    except:
        log("ERROR", "GET on {} failed".format(request_url))
        return None

def report_request(url, key, server):
    request_url = "http://{}/api/v3/{}/url/report/{}.json".format(
            server,
            key,
            url)

    print "Executing [GET {}]".format(request_url)
    try:
        response = requests.get(request_url)
        response.raise_for_status()
        return response.json()
    except:
        log("ERROR", "GET on {} failed".format(request_url))
        return None


def blacklisting_status_request(url, key, server):
    request_url = "http://{}/api/v3/{}/blacklist/status/{}.json".format(
            server,
            key,
            url)

    print "Executing [GET {}]".format(request_url)
    try:
        response = requests.get(request_url)
        response.raise_for_status()
        return response.json()
    except:
        log("ERROR", "GET on {} failed".format(request_url))
        return None


def is_done(url, key, server):
    st = status_request(url, key, server)
    if not st:
        log("ERROR", "Failed to receive reply from {}".format(server))
        return True
    if st["error"] != 200:
        log("ERROR", "Status request for {} returned {}".format(url, st))
        return True
    if st["status"]["state"].lower() != "done":
        return False
    return True


def scan_url(url, key, server):
    if not url or not key or not server:
        click.echo("Provided invalid input. Rerun with --help")
        return

    blst = blacklisting_status_request(url, key, server);
    click.echo(str(blst))

    st = scan_request(url, key, server)
    if not st:
        log("ERROR","Failed to scan URL: {}".format(url))
        click.echo("Failed to scan URL: {}".format(url))
        return

    log("INFO", "Output: " + str(st))

    if st["error"] != 200:
        log("ERROR",
            "Failed to submit scan request for {}. {} returned {}".format(
                url,
                server,
                st["errorstr"]))
    else:
        log("INFO", "Scan for {} submitted".format(url))

    while not is_done(url,key,server):
        time.sleep(2)
        log("DEBUG", "{} running".format(url))

    pprint(report_request(url,key,server))


@click.group()
@click.option('--verbose', '-v',
              is_flag=True,
              help='Enables verbose mode.',
              default=None)
@click.version_option('1.0')
@click.pass_context
def cli(ctx, verbose):
    # Create a repo object and remember it as as the context object.  From
    # this point onwards other commands can refer to it by using the
    # @pass_repo decorator.
    if verbose:
        ctx.obj.verbose = verbose


@cli.command(short_help='Perform scan of single URL')
@click.option('--url','-u',help="required URL",type = click.STRING)
@click.option('--key','-k',help="API key to use",type = click.STRING)
@click.option('--server','-s',help="API server name",type = click.STRING)
def scan(url,key,server):
    return scan_url(url, key, server)


if __name__ == '__main__':
    random.seed(time.time())
    cli();