import vt
import os
import requests
import virustotal_python
from dotenv import load_dotenv
from pprint import pprint
from base64 import urlsafe_b64encode

# todo: implement my own API request module to then try and cache the response (see -> https://realpython.com/caching-external-api-requests/#requests-cache)

def vt_lookup(URL):
    load_dotenv()
    api_key = os.getenv("VT_API")
    with virustotal_python.Virustotal(api_key) as vtotal:
        try:
            resp = vtotal.request("urls", data={"url": URL}, method="POST")
            print(resp)
            # Safe encode URL in base64 format
            # https://developers.virustotal.com/reference/url
            url_id = urlsafe_b64encode(URL.encode()).decode().strip("=")
            report = vtotal.request(f"urls/{url_id}")
            return report.data
        except virustotal_python.VirustotalError as err:
            print(f"Failed to send URL: {URL} for analysis and get the report: {err}")