import requests
import hashlib
import os
import time

# Update these two
API_KEY = ''
FILE_TO_UPLOAD = '' # sample.lime

API_BASE = 'https://api.deepbits.com/api/v1'

default_headers = {"x-api-key": API_KEY}

def upload_file(file_path, api_key):
    with open(file_path, "rb") as f:
        file_contents = f.read()

    print(f"Get upload URL for {os.path.basename(file_path)}...")
    file_name = os.path.basename(file_path)
    upload_url_response = requests.post(
        f"{API_BASE}/mdump/upload_url",
        json={"fileName": file_name},
        headers=default_headers
    )

    print(upload_url_response)

    if (upload_url_response.json()['meta']['code'] == 403):
        print("API key is invalid")
        raise Exception("API key is invalid")

    put_url = upload_url_response.json()['data']['uploadUrl']

    headers = {
        "Content-Type": "application/octet-stream",
        "x-api-key": api_key, 
        }

    print(f"Upload {os.path.basename(file_path)} to S3...")
    response = requests.put(
        put_url,
        data=file_contents,
        headers=headers
    )

    print(f"Trigger upload_success...")
    upload_success_response = requests.put(
        f"{API_BASE}/mdump/upload_success",
        json={
            "path": upload_url_response.json()['data']['path'],
            "hash": hashlib.sha256(file_contents).hexdigest(),
            "fileName": os.path.basename(file_path),
            "tags": ["tag1", "tag2"]
        },
        headers=default_headers
    )

    print(f"Upload Success Response: {upload_success_response.json()}\n")
    return upload_success_response.json()

def get_detail(detail_id):
    response = requests.get(
        f"{API_BASE}/mdump/detail/{detail_id}",
        headers=default_headers,
    )

    return response.json()


# upload file
upload_res = upload_file(FILE_TO_UPLOAD, API_KEY)
detail_id = upload_res['data']['_id']

# check scan status
detail_res = get_detail(detail_id)

print('sample detail')
print(detail_res)

print("Demo run successfully without error, get API docs on https://docs.deepbits.com/reference/scanmemorydumpid")