#!/usr/bin/env python
from __future__ import print_function
from json import dumps
from os import system
from tqdm import tqdm 
from hashlib import sha256
from rich import print_json
import os
from argparse import ArgumentParser
import traceback
import requests
import json
from os.path import isfile
from sys import exit

__author__ = "James Patrick-Evans"
__version__ = 0.01

re_conf = {
    'apikey' : 'l1br3', 
    'host' : 'https://api.reveng.ai'
}

def reveng_req(r: requests.request, end_point: str, data=None, ex_headers: dict = None, params=None):
    url = f"{re_conf['host']}/{end_point}"
    headers = { "Authorization": f"Bearer {re_conf['apikey']}" }
    if ex_headers:
        headers.update(ex_headers)
    return r(url, headers=headers, data=data, params=params)


def RE_delete(fpath: str):
    """
        Delete analysis results for Binary ID in command
    """
    bin_id = binary_id(fpath)
    res = reveng_req(requests.delete, f"{bin_id}")
    if res.status_code == 200:
        print(f"[+] Success. Securely deleted {fpath} analysis")
    elif res.status_code == 404:
        print("[!] Error, binary analysis not found.")
    else:
        print(f"[!] Error deleteing binary {bin_id}. Server returned {res.status_code}.")
    return


def RE_analyse(fpath: str):
    """
        Start analysis job for binary file
    """
    res = reveng_req(requests.post, f"analyse", data=open(fpath, 'rb').read())
    if res.status_code == 200:
        print("[+] Successfully submitted binary for analysis.")
        print(f"[+] {fpath} - {binary_id(fpath)}")
        return res

    if res.status_code == 400:
        if 'already exists' in json.loads(res.text)['reason']:
            print(f"[-] {fpath} already analysed. Please check the results log file for {binary_id(fpath)}")
            return True

    res.raise_for_status()


def RE_embeddings(fpath: str):
    """
        Fetch symbol embeddings
    """
    res = reveng_req(requests.get, f"embeddings/{binary_id(fpath)}")
    if res.status_code == 425:
        print(f"[-] Analysis for {binary_id(fpath)} still in progress. Please check the logs (-l) and try again later.")
        return

    res.raise_for_status()
    embeddings = res.json()
    print_json(data=embeddings)


def RE_logs(fpath: str):
    """
        Delete analysis results for Binary ID in command
    """
    bin_id = binary_id(fpath)
    res = reveng_req(requests.get, f"/log/{bin_id}")
    if res.status_code == 200:
        print(res.text)
        return
    elif res.status_code == 404:
        print(f"[!] Error, binary analysis for {bin_id} not found.")
        return

    res.raise_for_status()


def RE_nearest_symbols(embedding: list, nns: int = 5):
    """
        Get function name suggestions for an embedding
        :param nns: Number of nearest neighbors
    """
    res = reveng_req(requests.post, "ann", data=json.dumps(embedding), params={'nns': nns})
    res.raise_for_status()
    f_suggestions = res.json()
    print_json(data=f_suggestions)

def binary_id(path: str):
    """Take the SHA-256 hash of binary file"""
    hf = sha256()
    with open(path, "rb") as f:
        c = f.read()
        hf.update(c)
    return hf.hexdigest()


def parse_config():
    """
        Parse ~/.reveng.ai.yaml config file
    """     
    if not os.path.exists("~/.reveng.ai.yaml"):
        return

    with open("~/.reveng.ai.yaml", "r") as file:
        print("Parsing ~/.revengai.yaml")
        config = yaml.safe_load(file)
        for key in ('apikey', 'server'):
            if key in config:
                re_conf[key] = config[key]

if __name__ == '__main__':
    parse_config()
    parser = ArgumentParser()
    parser.add_argument("-b", "--binary", help="Path on binary to analyse")
    parser.add_argument("-a", "--analyse", action='store_true', help="Analyse new binary")
    parser.add_argument("-n", "--ann", action='store_true', help="Fetch Approximate Nearest Neighbours (ANNs) for embedding")
    parser.add_argument("-e", "--embedding", help="Path of JSON file containing a BinNet embedding")
    parser.add_argument("-m", "--model", default="BinNet", help="AI model used to generate embeddings")
    parser.add_argument("-x", "--extract", action='store_true', help="Fetch embeddings for binary")
    parser.add_argument("-l", "--logs", action='store_true', help="Fetch analysis log file for binary")
    parser.add_argument("-d", "--delete", action='store_true', help="Securely delete all analyses and metadata associated with binary")
    parser.add_argument("-k", "--api-key", default='l1br3', help="RevEng.AI API key")
    parser.add_argument("-s", "--host", default='https://api.reveng.ai', help="Analysis Host")
    parser.add_argument("--nns", default="5", help="Number of approximate nearest neighbors to fetch")

    args = parser.parse_args()
    if args.api_key:
        re_conf['api_key'] = args.api_key
    if args.host:
        re_conf['host'] = args.host

    if args.analyse or args.extract or args.logs or args.delete:
        # verify binary is a file
        if not os.path.isfile(args.binary):
            print("[!] Error, please supply a valid binary file using '-b'.")
            parser.print_help()
            exit(-1)

    if args.analyse:
        RE_analyse(args.binary)

    elif args.extract:
        RE_embeddings(args.binary)

    elif args.ann:
        # parse embedding json file
        if not isfile(args.embedding):
            print("[!] Error, please supply a valid embedding JSON file using '-e'")
            parser.print_help()
            exit(-1)

        embedding = json.loads(open(args.embedding, 'r').read())
        RE_nearest_symbols(embedding, int(args.nns))

    elif args.logs:
        RE_logs(args.binary)

    elif args.delete:
        RE_delete(args.binary)

    else:
        print("[!] Error, please suply an action command")
        parser.print_help()

