#!/usr/bin/env python3

'''
This is a simple script meant to sumbmit and retrieve 1 DARKINT Score from the DarkOwl Vision Score API
The input criteria is in the 'payload' variable in the main function
All results will be saved in a folder called 'Results' in the current directory
If this folder does not exist, it will be created
'''
import requests
import json
import base64
import datetime
import hashlib
import hmac
import os
from datetime import datetime
from time import sleep, time

publicKey = 'PUBLIC_KEY_HERE'
privateKey = 'PRIVATE_KEY_HERE'

host = 'api.darkowl.com'

def main():

    payload = {'domains': ['test.com'], 'emailDomains': ['test.com', '*.test.com']}

    ID = submit(payload)
    counter = 0
    sleepVal = 1
    while True:
        sleep(sleepVal)
        if status(ID) == 'completed':
            result(ID)
            break
        if counter >= 3 and counter < 5:
            sleepVal = 2
        else:
            sleepVal = 5
        counter += 1


def submit(payload):
    headers = gen_headers('POST', 'submit')

    payload_json = json.dumps(payload)

    r = requests.post(f'https://{host}/api/v1/score/submit', headers=headers, data=payload_json)
    if r.status_code != 202:
        print(r.content)
        print(payload)
        exit(1)

    print(f'Submission Success, ID {r.text}')
    return r.text

def status(ID):
    headers = gen_headers('GET', f'status?id={ID}')

    r = requests.get(f'https://{host}/api/v1/score/status?id={ID}', headers=headers)
    if r.status_code != 200:
        print(r.content)
        print(r.status_code)
        exit(1)
    resp_json = r.json()
    status = resp_json['code']
    progress = resp_json['text']
    print(f'Status: {status} ... {progress}')

    return status

def result(ID):
    headers = gen_headers('GET', f'result?id={ID}')

    r = requests.get(f'https://{host}/api/v1/score/result?id={ID}', headers=headers)
    if r.status_code != 200:
        print(r.content)
        exit(1)

    dump_result_to_file(r.json())

def dump_result_to_file(resp_json):
    print('\nRESULT:')
    print(json.dumps(resp_json, indent=4))
    date = str(resp_json['result']['date'])

    if not os.path.exists('./Results/'):
        os.mkdir('./Results/')

    with open(f'./Results/{date}.json', 'w') as outFile:
        json.dump(resp_json, outFile, indent=4)


def gen_headers(verb, endpoint):

    path = '/api/v1/score/'

    date = datetime.utcnow().strftime('%a, %d %b %Y %H:%M:%S GMT')
    auth = generate_auth_header(path + endpoint, verb, privateKey, publicKey, date)
    headers = {'Authorization': auth, 'X-VISION-DATE': date, 'content-type': 'application/json'}
    return headers

def generate_auth_header(abs_path,
                         http_method,
                         private_key,
                         public_key,
                         time_stamp):

    string2hash = http_method + abs_path + time_stamp
    bkey = bytes(source=private_key, encoding='UTF-8')
    bpayload = bytes(source=string2hash, encoding='UTF-8')
    hmacsha1 = hmac.new(bkey, bpayload, hashlib.sha1).digest()
    base64encoded = base64.b64encode(hmacsha1).decode('UTF-8')
    auth_header = f'OWL {public_key}:{base64encoded}'
    return auth_header


if __name__ == '__main__':
    main()