vt.py

#!/usr/bin/env python3

import argparse
import base64
import getpass
import hashlib
import json
import os
import sys
import time
import urllib.parse
from datetime import datetime
from pathlib import Path

import requests


VT_API = "https://www.virustotal.com/api/v3"
VT_WEB = "https://www.virustotal.com"

API_KEY_URL = f"{VT_WEB}/gui/my-apikey"

CONFIG_DIR = Path.home() / ".config" / "vt"
CONFIG_FILE = CONFIG_DIR / "config"

REQUEST_TIMEOUT = 30
UPLOAD_TIMEOUT = 600
POLL_INTERVAL = 5

MAX_DIRECT_UPLOAD_SIZE = 32 * 1024 * 1024
MAX_UPLOAD_SIZE = 650 * 1024 * 1024


class VirusTotalError(Exception):
    """Base exception for VirusTotal errors."""

    def __init__(self, message, code=None, status=None):
        super().__init__(message)
        self.code = code
        self.status = status


class VirusTotalQuotaError(VirusTotalError):
    """Raised when a VirusTotal quota or rate limit is reached."""


class VirusTotal:
    def __init__(self, api_key):
        self.api_key = api_key

        self.session = requests.Session()
        self.session.headers.update({
            "x-apikey": api_key,
            "accept": "application/json",
        })

    # ------------------------------------------------------------------
    # API
    # ------------------------------------------------------------------

    def request(self, method, endpoint, **kwargs):
        """Make an authenticated VirusTotal API request."""

        if endpoint.startswith(("http://", "https://")):
            url = endpoint
        else:
            url = f"{VT_API}{endpoint}"

        try:
            response = self.session.request(
                method,
                url,
                timeout=REQUEST_TIMEOUT,
                **kwargs,
            )

        except requests.RequestException as error:
            raise VirusTotalError(
                f"Network error: {error}"
            ) from error

        if response.ok:
            if not response.content:
                return None

            try:
                return response.json()

            except ValueError as error:
                raise VirusTotalError(
                    "VirusTotal returned invalid JSON."
                ) from error

        try:
            error_data = response.json()
            error = error_data.get("error", {})

            error_code = error.get("code")
            error_message = error.get(
                "message",
                response.text,
            )

        except (ValueError, AttributeError):
            error_code = None
            error_message = response.text

        if error_code in (
            "QuotaExceededError",
            "TooManyRequestsError",
        ):
            raise VirusTotalQuotaError(
                error_message,
                code=error_code,
                status=response.status_code,
            )

        raise VirusTotalError(
            error_message,
            code=error_code,
            status=response.status_code,
        )

    # ------------------------------------------------------------------
    # Authentication
    # ------------------------------------------------------------------

    def verify_api_key(self):
        """Verify the configured VirusTotal API key."""

        return self.request(
            "GET",
            f"/users/{self.api_key}",
        )

    # ------------------------------------------------------------------
    # User / quota
    # ------------------------------------------------------------------

    def get_user(self):
        """Retrieve the authenticated user's information."""

        return self.request(
            "GET",
            f"/users/{self.api_key}",
        )

    # ------------------------------------------------------------------
    # Files
    # ------------------------------------------------------------------

    def get_file(self, sha256):
        """Retrieve a file by SHA-256."""

        return self.request(
            "GET",
            f"/files/{sha256}",
        )

    def get_upload_url(self):
        """Obtain a dedicated VirusTotal upload URL."""

        result = self.request(
            "GET",
            "/files/upload_url",
        )

        return result["data"]

    def upload_file(self, path):
        """
        Upload a file.

        Files <= 32 MiB use POST /files.
        Larger files use the dedicated upload URL.
        """

        file_size = path.stat().st_size

        if file_size > MAX_UPLOAD_SIZE:
            raise VirusTotalError(
                "File exceeds VirusTotal's 650 MB upload limit."
            )

        if file_size <= MAX_DIRECT_UPLOAD_SIZE:
            return self._upload_small_file(path)

        return self._upload_large_file(path)

    def _upload_small_file(self, path):
        """Upload a file of 32 MiB or less."""

        try:
            with path.open("rb") as file:
                response = self.session.post(
                    f"{VT_API}/files",
                    files={
                        "file": (
                            path.name,
                            file,
                            "application/octet-stream",
                        )
                    },
                    timeout=UPLOAD_TIMEOUT,
                )

        except requests.RequestException as error:
            raise VirusTotalError(
                f"File upload failed: {error}"
            ) from error

        return self._process_upload_response(response)

    def _upload_large_file(self, path):
        """Upload a file larger than 32 MiB."""

        upload_url = self.get_upload_url()

        try:
            with path.open("rb") as file:
                response = requests.post(
                    upload_url,
                    files={
                        "file": (
                            path.name,
                            file,
                            "application/octet-stream",
                        )
                    },
                    timeout=UPLOAD_TIMEOUT,
                )

        except requests.RequestException as error:
            raise VirusTotalError(
                f"Large file upload failed: {error}"
            ) from error

        return self._process_upload_response(response)

    @staticmethod
    def _process_upload_response(response):
        """Process a VirusTotal upload response."""

        if not response.ok:
            try:
                data = response.json()
                error = data.get("error", {})

                message = error.get(
                    "message",
                    response.text,
                )

                code = error.get("code")

            except (ValueError, AttributeError):
                message = response.text
                code = None

            raise VirusTotalError(
                f"File upload failed: {message}",
                code=code,
                status=response.status_code,
            )

        try:
            return response.json()

        except ValueError as error:
            raise VirusTotalError(
                "VirusTotal returned an invalid upload response."
            ) from error

    # ------------------------------------------------------------------
    # URLs
    # ------------------------------------------------------------------

    def submit_url(self, url):
        """Submit a URL to VirusTotal."""

        return self.request(
            "POST",
            "/urls",
            data={
                "url": url,
            },
        )

    def get_url(self, url_id):
        """Retrieve a VirusTotal URL object."""

        return self.request(
            "GET",
            f"/urls/{url_id}",
        )

    def scan_url(self, url):
        """Submit and analyse a URL."""

        result = self.submit_url(url)

        analysis_id = (
            result
            .get("data", {})
            .get("id")
        )

        if not analysis_id:
            raise VirusTotalError(
                "VirusTotal URL submission did not "
                "return an analysis ID."
            )

        submitted_at = current_timestamp()

        analysis = self.wait_for_analysis(
            analysis_id,
            submitted_at,
            "URL Submitted to Virus Total",
        )

        url_id = get_url_id(url)

        report = self.get_url(url_id)

        return {
            "type": "url",
            "target": url,
            "report": report,
            "analysis": analysis,
            "report_url": get_url_report_url(url),
        }

    # ------------------------------------------------------------------
    # Analysis
    # ------------------------------------------------------------------

    def get_analysis(self, analysis_id):
        """Retrieve an analysis object."""

        return self.request(
            "GET",
            f"/analyses/{analysis_id}",
        )

    def wait_for_analysis(
        self,
        analysis_id,
        initial_timestamp,
        initial_status,
    ):
        """
        Poll an analysis until completion.

        A status is recorded only when it changes.
        """

        print()
        print("Status:")
        print(
            f"[{initial_timestamp}] - "
            f"{initial_status}"
        )

        previous_status = None

        while True:
            result = self.get_analysis(
                analysis_id
            )

            status = (
                result
                .get("data", {})
                .get("attributes", {})
                .get("status")
            )

            if status != previous_status:
                timestamp = current_timestamp()

                display_status = {
                    "queued": "Queued for Analysis",
                    "in-progress": "Analysis In Progress",
                    "completed": "Analysis Complete",
                }.get(
                    status,
                    f"Status: {status}",
                )

                print(
                    f"[{timestamp}] - "
                    f"{display_status}"
                )

                previous_status = status

            if status == "completed":
                return result

            time.sleep(POLL_INTERVAL)

    # ------------------------------------------------------------------
    # File scanning
    # ------------------------------------------------------------------

    def scan_file(self, path):
        """Look up a file and upload it if necessary."""

        sha256 = calculate_sha256(path)

        print(
            f"SHA256: {sha256}"
        )

        try:
            report = self.get_file(
                sha256
            )

            print()
            print(
                "File already exists in VirusTotal."
            )

            return {
                "type": "file",
                "target": str(path),
                "sha256": sha256,
                "report": report,
                "report_url": get_file_report_url(
                    sha256
                ),
            }

        except VirusTotalError as error:
            if error.status != 404:
                raise

        print()
        print(
            "File not found in VirusTotal."
        )

        print(
            "Uploading..."
        )

        upload_result = self.upload_file(
            path
        )

        uploaded_at = current_timestamp()

        analysis_id = (
            upload_result
            .get("data", {})
            .get("id")
        )

        if not analysis_id:
            raise VirusTotalError(
                "VirusTotal upload response did not "
                "contain an analysis ID."
            )

        analysis = self.wait_for_analysis(
            analysis_id,
            uploaded_at,
            "File Uploaded to Virus Total",
        )

        print()
        print(
            "Retrieving final file report..."
        )

        report = self.get_file(
            sha256
        )

        return {
            "type": "file",
            "target": str(path),
            "sha256": sha256,
            "report": report,
            "analysis": analysis,
            "report_url": get_file_report_url(
                sha256
            ),
        }


# ----------------------------------------------------------------------
# Configuration
# ----------------------------------------------------------------------


def load_api_key():
    """Load the API key from ~/.config/vt/config."""

    if not CONFIG_FILE.exists():
        return None

    try:
        with CONFIG_FILE.open(
            "r",
            encoding="utf-8",
        ) as file:

            for line in file:
                line = line.strip()

                if line.startswith("api_key="):
                    return line.split(
                        "=",
                        1,
                    )[1].strip()

    except OSError as error:
        raise VirusTotalError(
            f"Unable to read configuration: {error}"
        ) from error

    return None


def save_api_key(api_key):
    """Save the API key with restrictive permissions."""

    CONFIG_DIR.mkdir(
        parents=True,
        exist_ok=True,
        mode=0o700,
    )

    old_umask = os.umask(0o077)

    try:
        with CONFIG_FILE.open(
            "w",
            encoding="utf-8",
        ) as file:

            file.write(
                f"api_key={api_key}\n"
            )

    finally:
        os.umask(old_umask)

    try:
        CONFIG_DIR.chmod(0o700)
        CONFIG_FILE.chmod(0o600)

    except OSError:
        pass


def setup_api_key():
    """Perform first-run API key setup."""

    print()
    print("VirusTotal API key is not configured.")
    print()
    print("Open the following URL to obtain your API key:")
    print()
    print(f"  {API_KEY_URL}")
    print()

    print("Enter your VirusTotal API key.")
    print("The key will not be displayed.")
    print()

    api_key = getpass.getpass(
        "API key: "
    ).strip()

    if not api_key:
        raise VirusTotalError(
            "No API key was entered."
        )

    print()
    print("Verifying API key...")

    vt = VirusTotal(api_key)

    try:
        vt.verify_api_key()

    except VirusTotalError as error:
        raise VirusTotalError(
            f"API key verification failed: {error}"
        ) from error

    save_api_key(api_key)

    print("API key verified successfully.")
    print(
        f"Configuration saved to: {CONFIG_FILE}"
    )
    print()

    return api_key


def get_api_key():
    """Load an existing API key or perform setup."""

    api_key = load_api_key()

    if api_key:
        return api_key

    return setup_api_key()


# ----------------------------------------------------------------------
# Utility functions
# ----------------------------------------------------------------------


def current_timestamp():
    """Return the current local timestamp."""

    return datetime.now().strftime(
        "%Y-%b-%d %H:%M:%S"
    )


def calculate_sha256(path):
    """Calculate a file's SHA-256 hash."""

    digest = hashlib.sha256()

    with path.open("rb") as file:
        while True:
            chunk = file.read(
                1024 * 1024
            )

            if not chunk:
                break

            digest.update(chunk)

    return digest.hexdigest()


def is_url(value):
    """Return True if value is an HTTP(S) URL."""

    parsed = urllib.parse.urlparse(
        value
    )

    return (
        parsed.scheme.lower()
        in ("http", "https")
        and bool(parsed.netloc)
    )


def get_url_id(url):
    """
    Generate VirusTotal's URL object ID.

    VirusTotal uses URL-safe Base64 without
    trailing '=' padding.
    """

    encoded = (
        base64.urlsafe_b64encode(
            url.encode("utf-8")
        )
        .decode("ascii")
    )

    return encoded.rstrip("=")


def get_file_report_url(sha256):
    """Return the VirusTotal file report URL."""

    return (
        f"{VT_WEB}/gui/file/"
        f"{sha256}/detection"
    )


def get_url_report_url(url):
    """Return the VirusTotal URL report URL."""

    return (
        f"{VT_WEB}/gui/url/"
        f"{get_url_id(url)}/detection"
    )


# ----------------------------------------------------------------------
# Results
# ----------------------------------------------------------------------


def get_stats(report):
    """Extract analysis statistics from a VT object."""

    attributes = (
        report
        .get("data", {})
        .get("attributes", {})
    )

    return attributes.get(
        "last_analysis_stats",
        attributes.get(
            "stats",
            {},
        ),
    )


def get_verdict(stats):
    """Generate a simple scan verdict."""

    if stats.get(
        "malicious",
        0,
    ) > 0:
        return "MALICIOUS"

    if stats.get(
        "suspicious",
        0,
    ) > 0:
        return "SUSPICIOUS"

    return "CLEAN"


def print_stats(stats):
    """Print scan statistics."""

    print()
    print("VirusTotal")
    print("───────────")

    print(
        f"Malicious : "
        f"{stats.get('malicious', 0)}"
    )

    print(
        f"Suspicious: "
        f"{stats.get('suspicious', 0)}"
    )

    print(
        f"Harmless  : "
        f"{stats.get('harmless', 0)}"
    )

    print(
        f"Undetected: "
        f"{stats.get('undetected', 0)}"
    )

    print(
        f"Timeout   : "
        f"{stats.get('timeout', 0)}"
    )


def print_result(result):
    """Print a human-readable scan result."""

    stats = get_stats(
        result["report"]
    )

    verdict = get_verdict(
        stats
    )

    print_stats(
        stats
    )

    print()
    print(
        f"Verdict: {verdict}"
    )

    print()
    print(
        "Report:"
    )

    print(
        result["report_url"]
    )


# ----------------------------------------------------------------------
# Quota
# ----------------------------------------------------------------------


def print_quota(vt):
    """Display the user's VirusTotal API quota."""

    user = vt.get_user()

    attributes = (
        user
        .get("data", {})
        .get("attributes", {})
    )

    quotas = attributes.get(
        "quotas",
        {},
    )

    daily = quotas.get(
        "api_requests_daily"
    )

    if daily is None:
        print(
            "VirusTotal did not provide "
            "daily API quota information."
        )

        return

    allowed = daily.get(
        "allowed",
        0,
    )

    used = daily.get(
        "used",
        0,
    )

    remaining = max(
        allowed - used,
        0,
    )

    print()
    print(
        "VirusTotal API Usage"
    )
    print(
        "────────────────────"
    )

    print()
    print(
        "Daily"
    )

    print(
        f"  Used:       {used}"
    )

    print(
        f"  Remaining:  {remaining}"
    )

    print(
        f"  Limit:      {allowed}"
    )

    print()
    print(
        "Rate"
    )

    print(
        "  Limit:      4 requests/minute"
    )

    print()
    print(
        "Daily reset"
    )

    print(
        "  00:00 UTC"
    )


# ----------------------------------------------------------------------
# JSON
# ----------------------------------------------------------------------


def print_json(result):
    """Print a machine-readable scan result."""

    stats = get_stats(
        result["report"]
    )

    output = {
        "type": result["type"],
        "target": result["target"],
        "malicious": stats.get(
            "malicious",
            0,
        ),
        "suspicious": stats.get(
            "suspicious",
            0,
        ),
        "harmless": stats.get(
            "harmless",
            0,
        ),
        "undetected": stats.get(
            "undetected",
            0,
        ),
        "timeout": stats.get(
            "timeout",
            0,
        ),
        "verdict": get_verdict(
            stats
        ),
        "report_url": result[
            "report_url"
        ],
    }

    if result["type"] == "file":
        output["sha256"] = result[
            "sha256"
        ]

    print(
        json.dumps(
            output,
            indent=2,
        )
    )


# ----------------------------------------------------------------------
# Target handling
# ----------------------------------------------------------------------


def scan_target(vt, target):
    """Scan a URL or local file."""

    if is_url(target):
        print(
            f"URL: {target}"
        )

        return vt.scan_url(
            target
        )

    path = Path(
        target
    ).expanduser()

    if not path.is_file():
        raise VirusTotalError(
            f"File does not exist: {target}"
        )

    print(
        f"File: {path}"
    )

    return vt.scan_file(
        path
    )


# ----------------------------------------------------------------------
# CLI
# ----------------------------------------------------------------------


def main():
    parser = argparse.ArgumentParser(
        prog="vt",
        description=(
            "Scan files and URLs using "
            "the VirusTotal API."
        ),
    )

    parser.add_argument(
        "target",
        nargs="?",
        help="File or URL to scan",
    )

    parser.add_argument(
        "--quota",
        action="store_true",
        help="Display VirusTotal API quota",
    )

    parser.add_argument(
        "--json",
        action="store_true",
        help="Output scan result as JSON",
    )

    args = parser.parse_args()

    if args.quota and args.target:
        parser.error(
            "--quota cannot be used with a target"
        )

    if args.json and args.quota:
        parser.error(
            "--json cannot be used with --quota"
        )

    if not args.target and not args.quota:
        parser.print_help()
        return 0

    try:
        api_key = get_api_key()

        vt = VirusTotal(
            api_key
        )

        if args.quota:
            print_quota(
                vt
            )

            return 0

        result = scan_target(
            vt,
            args.target,
        )

        if args.json:
            print_json(
                result
            )

        else:
            print_result(
                result
            )

        return 0

    except VirusTotalQuotaError as error:
        print(
            file=sys.stderr
        )

        print(
            "VirusTotal API limit reached.",
            file=sys.stderr,
        )

        if error.code:
            print(
                f"Reason: {error.code}",
                file=sys.stderr,
            )

        print(
            error,
            file=sys.stderr,
        )

        if error.code == "QuotaExceededError":
            print(
                file=sys.stderr
            )

            print(
                "Daily quotas reset at 00:00 UTC.",
                file=sys.stderr,
            )

        elif error.code == "TooManyRequestsError":
            print(
                file=sys.stderr
            )

            print(
                "Please wait before trying again.",
                file=sys.stderr,
            )

        return 2

    except VirusTotalError as error:
        print(
            f"Error: {error}",
            file=sys.stderr,
        )

        if error.code:
            print(
                f"VirusTotal error code: "
                f"{error.code}",
                file=sys.stderr,
            )

        return 1

    except KeyboardInterrupt:
        print(
            "\nOperation cancelled.",
            file=sys.stderr,
        )

        return 130


if __name__ == "__main__":
    sys.exit(
        main()
    )