#!/usr/bin/env python3 """Configure persistent VirusTotal access for one local MCP client. Python 3.11+ on Windows, macOS or Linux. Run --help for options. Credentials stay in protected local files; registration is performed at most once per setup attempt. """ import argparse import base64 import contextlib import hashlib import http.client import io import json import os from pathlib import Path import re import secrets import shlex import shutil import stat import subprocess import sys import tempfile if sys.version_info < (3, 11): raise SystemExit("VirusTotal setup requires Python 3.11 or newer.") import tomllib HOST = "ai.virustotal.com" MCP_URL = f"https://{HOST}/mcp" PACKAGE_VERSION = "0.9.5" PACKAGE = f"vt-mcp=={PACKAGE_VERSION}" MAX_CONFIG_BYTES = 4 * 1024 * 1024 MAX_RESPONSE_BYTES = 16 * 1024 TOKEN_PATTERN = re.compile(r"vtai_[A-Za-z0-9_-]{1,507}\Z") MANAGED_HELPER_SHA256 = frozenset({ # Previously distributed installers; never overwrite unknown local code. "d0090a1fd30bb680b930bfe7c5a73c320f4a43b02703c2a8642244a70a1ca955", "8e02c46fd2003058cc443f1d2131a1b82405b6600004c7d9e717a83deefdaaad", "90fcac59da00aec5983157ebc40a71322f276a1d44327a21b92cdb2a08b7ee43", }) AUTH_FIELDS = {"headers", "http_headers", "env_http_headers", "bearer_token_env_var", "headersHelper", "http_headers_helper"} class SetupError(Exception): """An actionable error that never contains a credential or response body.""" def __init__(self, message, code="setup_error"): super().__init__(message) self.code = code class RegistrationLimited(SetupError): """A definite quota rejection, before any identity was committed.""" class WindowsFiles: """Small Win32 file adapter. chmod is not a credential ACL on Windows. New objects have a protected DACL granting only the current user access. Existing nonsecret files may also trust SYSTEM and Administrators, which already control the machine. Credential files require the current user only. """ def __init__(self): import ctypes as c from ctypes import wintypes as w self.c, self.w = c, w self.kernel = c.WinDLL("kernel32", use_last_error=True) self.advapi = c.WinDLL("advapi32", use_last_error=True) ptr = c.c_void_p signatures = [ (self.kernel, "GetCurrentProcess", w.HANDLE, []), (self.kernel, "GetSystemDirectoryW", w.UINT, [w.LPWSTR, w.UINT]), (self.kernel, "CloseHandle", w.BOOL, [w.HANDLE]), (self.kernel, "LocalFree", ptr, [ptr]), (self.kernel, "CreateFileW", w.HANDLE, [w.LPCWSTR, w.DWORD, w.DWORD, ptr, w.DWORD, w.DWORD, w.HANDLE]), (self.kernel, "CreateDirectoryW", w.BOOL, [w.LPCWSTR, ptr]), (self.kernel, "GetFileInformationByHandle", w.BOOL, [w.HANDLE, ptr]), (self.kernel, "GetFileType", w.DWORD, [w.HANDLE]), (self.kernel, "MoveFileExW", w.BOOL, [w.LPCWSTR, w.LPCWSTR, w.DWORD]), (self.advapi, "OpenProcessToken", w.BOOL, [w.HANDLE, w.DWORD, ptr]), (self.advapi, "GetTokenInformation", w.BOOL, [w.HANDLE, c.c_int, ptr, w.DWORD, ptr]), (self.advapi, "ConvertSidToStringSidW", w.BOOL, [ptr, ptr]), (self.advapi, "ConvertStringSecurityDescriptorToSecurityDescriptorW", w.BOOL, [w.LPCWSTR, w.DWORD, ptr, ptr]), (self.advapi, "GetSecurityInfo", w.DWORD, [w.HANDLE, c.c_int, w.DWORD, ptr, ptr, ptr, ptr, ptr]), (self.advapi, "GetAce", w.BOOL, [ptr, w.DWORD, ptr]), ] for dll, name, result, arguments in signatures: function = getattr(dll, name) function.restype, function.argtypes = result, arguments class SecurityAttributes(c.Structure): _fields_ = [("length", w.DWORD), ("descriptor", ptr), ("inherit", w.BOOL)] class FileInformation(c.Structure): _fields_ = [("attributes", w.DWORD), ("created", w.FILETIME), ("accessed", w.FILETIME), ("written", w.FILETIME), ("volume", w.DWORD), ("size_high", w.DWORD), ("size_low", w.DWORD), ("links", w.DWORD), ("index_high", w.DWORD), ("index_low", w.DWORD)] class ACL(c.Structure): _fields_ = [("revision", w.BYTE), ("reserved", w.BYTE), ("size", w.WORD), ("count", w.WORD), ("reserved2", w.WORD)] self.Attributes, self.Information, self.ACL = SecurityAttributes, FileInformation, ACL token, size = w.HANDLE(), w.DWORD() self.require(self.advapi.OpenProcessToken(self.kernel.GetCurrentProcess(), 0x0008, c.byref(token))) try: self.advapi.GetTokenInformation(token, 1, None, 0, c.byref(size)) self.require(0 < size.value <= 16384) data = c.create_string_buffer(size.value) self.require(self.advapi.GetTokenInformation(token, 1, data, size, c.byref(size))) self.sid = self.sid_text(c.cast(data, c.POINTER(ptr))[0]) finally: self.kernel.CloseHandle(token) @staticmethod def require(ok): if not ok: raise SetupError("Windows could not verify private local storage. Use a local NTFS directory owned by your account.") def sid_text(self, sid): value = self.w.LPWSTR() self.require(self.advapi.ConvertSidToStringSidW(sid, self.c.byref(value))) try: return value.value finally: self.kernel.LocalFree(value) @contextlib.contextmanager def attributes(self, directory=False): descriptor = self.c.c_void_p() inheritance = "OICI" if directory else "" self.require(self.advapi.ConvertStringSecurityDescriptorToSecurityDescriptorW( f"O:{self.sid}D:P(A;{inheritance};FA;;;{self.sid})", 1, self.c.byref(descriptor), None)) try: yield self.Attributes(self.c.sizeof(self.Attributes), descriptor, False) finally: self.kernel.LocalFree(descriptor) def path(self, path): path = Path(path).absolute() if not re.fullmatch(r"[A-Za-z]:", path.drive) or any( ":" in part or part.endswith((".", " ")) or Path(part).is_reserved() for part in path.parts[1:] ): raise SetupError("Use a regular local Windows path, without device names, alternate streams or network shares.") for parent in (path, *path.parents): try: if parent.lstat().st_file_attributes & 0x400: raise SetupError("A setup path contains a Windows reparse point. Use a regular local directory.") except FileNotFoundError: continue return "\\\\?\\" + str(path) def check(self, handle, *, directory=False, private=False): c, w = self.c, self.w info = self.Information() self.require(self.kernel.GetFileInformationByHandle(handle, c.byref(info))) if (self.kernel.GetFileType(handle) != 1 or info.attributes & 0x400 or bool(info.attributes & 0x10) != directory or not directory and info.links != 1): raise SetupError("Setup paths must be regular files or directories, without reparse points or hard links.") owner, acl, descriptor = c.c_void_p(), c.c_void_p(), c.c_void_p() self.require(self.advapi.GetSecurityInfo(handle, 1, 0x0001 | 0x0004, c.byref(owner), None, c.byref(acl), None, c.byref(descriptor)) == 0) try: trusted = {self.sid} if private else {self.sid, "S-1-5-18", "S-1-5-32-544"} if self.sid_text(owner) not in trusted or not acl: raise SetupError("A setup file has an unsafe owner or ACL. Protect it before continuing.") count = c.cast(acl, c.POINTER(self.ACL)).contents.count self.require(count <= 4096) for index in range(count): ace = c.c_void_p() self.require(self.advapi.GetAce(acl, index, c.byref(ace))) header = (w.BYTE * 4).from_address(ace.value) if header[0] == 1 or header[1] & 0x08: # Deny or inheritance-only ACE. continue self.require(header[0] == 0) # Do not guess unfamiliar allow-ACE semantics. mask = w.DWORD.from_address(ace.value + 4).value principal = self.sid_text(ace.value + 8) # Nonsecret objects may be readable, but never writable, by other users. dangerous = mask if private else mask & 0x500D0156 if principal not in trusted and dangerous: raise SetupError("A setup file allows access by other users. Restrict its Windows Security permissions before continuing.") finally: self.kernel.LocalFree(descriptor) def open(self, path, *, create=False, exclusive=False, writable=False): import msvcrt access = 0x80000000 | (0x40000000 if writable else 0) with self.attributes() as attributes: handle = self.kernel.CreateFileW(self.path(path), access, 0 if exclusive else 1, self.c.byref(attributes), 4 if exclusive else 1 if create else 3, 0x00200000, None) error = self.c.get_last_error() if handle == self.c.c_void_p(-1).value: if error in (2, 3): raise FileNotFoundError() if error in (80, 183): raise FileExistsError() if exclusive and error == 32: raise SetupError("Another VirusTotal setup is running. Wait for it to finish and rerun.") raise SetupError("Cannot safely open a Windows setup file. Close programs editing it and check its permissions.") try: self.check(handle, private=create or exclusive) return msvcrt.open_osfhandle(handle, os.O_BINARY | (os.O_RDWR if writable else os.O_RDONLY)) except BaseException: self.kernel.CloseHandle(handle) raise def directory(self, path): name = self.path(path) if not path.exists(): private_directory(path.parent) with self.attributes(directory=True) as attributes: self.require(self.kernel.CreateDirectoryW(name, self.c.byref(attributes))) handle = self.kernel.CreateFileW(name, 0x00020080, 1, None, 3, 0x02200000, None) self.require(handle != self.c.c_void_p(-1).value) try: self.check(handle, directory=True) finally: self.kernel.CloseHandle(handle) def publish(self, source, destination, replace): if not self.kernel.MoveFileExW(self.path(source), self.path(destination), 0x8 | int(replace)): raise SetupError("Configuration changed or is open in another program. Close settings editors and rerun safely.") def system_powershell(self): directory = self.c.create_unicode_buffer(32768) size = self.kernel.GetSystemDirectoryW(directory, len(directory)) self.require(0 < size < len(directory)) executable = Path(directory.value) / "WindowsPowerShell" / "v1.0" / "powershell.exe" if not executable.is_file() or any(char in str(executable) for char in '%!"\r\n'): raise SetupError("Windows HTTP setup requires the system Windows PowerShell. Restore that component or use --transport stdio.") return str(executable) _windows_files = None def windows_files(): global _windows_files if _windows_files is None: _windows_files = WindowsFiles() return _windows_files def private_directory(path): """Create private directories without changing existing unrelated parents.""" path = Path(path) if os.name == "nt": return windows_files().directory(path) if any(parent.is_symlink() for parent in (path, *path.parents)): raise SetupError("A setup directory is a symbolic link; use a regular private directory.") if not path.exists(): private_directory(path.parent) path.mkdir(mode=0o700) if not path.is_dir(): raise SetupError("A required setup directory is not a directory.") if path.stat().st_uid != os.getuid(): raise SetupError("A setup directory is owned by another user.") if path.stat().st_mode & 0o022: raise SetupError("A setup directory is writable by other users; protect it before continuing.") def read_file(path, limit=MAX_CONFIG_BYTES, private=False): try: fd = windows_files().open(path) if os.name == "nt" else os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK) except FileNotFoundError: return None except OSError: raise SetupError("Cannot safely read a setup file; check its type and permissions.") from None with os.fdopen(fd, "rb") as source: info = os.fstat(source.fileno()) if os.name == "nt": import msvcrt windows_files().check(msvcrt.get_osfhandle(source.fileno()), private=private) elif not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid() or info.st_nlink != 1: raise SetupError("Setup files must be regular files owned by you, without hard links.") if os.name != "nt" and private and info.st_mode & 0o077: raise SetupError("The token file must have owner-only permissions (chmod 600).") value = source.read(limit + 1) if len(value) > limit: raise SetupError("A setup file exceeds the supported size limit.") return value def atomic_write(path, content, expected): """Protect new bytes before publishing; refuse a changed destination.""" private_directory(path.parent) if os.name == "nt": temporary = path.parent / (".vt-mcp-" + secrets.token_hex(12)) fd = windows_files().open(temporary, create=True, writable=True) else: fd, temporary = tempfile.mkstemp(prefix=".vt-mcp-", dir=path.parent) try: with os.fdopen(fd, "wb") as output: if os.name != "nt": os.fchmod(output.fileno(), 0o600) output.write(content) output.flush() os.fsync(output.fileno()) if read_file(path) != expected: raise SetupError("Configuration changed during setup. Close other settings editors and rerun.") if os.name == "nt": windows_files().publish(temporary, path, expected is not None) elif expected is None: # link is an atomic create-if-absent; never replace a newly created token. try: os.link(temporary, path) except FileExistsError: raise SetupError("A setup file was created concurrently. Rerun to reuse it.") from None else: os.replace(temporary, path) if os.name != "nt": directory_fd = os.open(path.parent, os.O_RDONLY) try: os.fsync(directory_fd) finally: os.close(directory_fd) finally: with contextlib.suppress(FileNotFoundError): os.unlink(temporary) @contextlib.contextmanager def setup_lock(directory): private_directory(directory) if os.name == "nt": fd = windows_files().open(directory / "setup.lock", exclusive=True, writable=True) try: yield finally: os.close(fd) return import fcntl fd = os.open(directory / "setup.lock", os.O_CREAT | os.O_RDWR | os.O_NOFOLLOW, 0o600) try: info = os.fstat(fd) if not stat.S_ISREG(info.st_mode) or info.st_uid != os.getuid() or info.st_nlink != 1: raise SetupError("The setup lock must be a regular file owned by you.") try: fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) except BlockingIOError: raise SetupError("Another VirusTotal setup is running. Wait for it to finish and rerun.") from None yield finally: os.close(fd) def valid_token(raw): try: token = raw.decode("ascii").strip() except (UnicodeError, AttributeError): raise SetupError("The token file is invalid. Import a valid VTAI token; do not register again.") from None if not TOKEN_PATTERN.fullmatch(token): raise SetupError("The token file is invalid. Import a valid VTAI token; do not register again.") return token def token_from_file(path): raw = read_file(path, 1024, private=True) return None if raw is None else valid_token(raw) def unique_object(pairs): result = {} for key, value in pairs: if key in result: raise ValueError("duplicate key") result[key] = value return result def parse_config(raw, client): try: text = (raw or b"").decode("utf-8") data = tomllib.loads(text) if client == "codex" else json.loads(text or "{}", object_pairs_hook=unique_object) key = "mcp_servers" if client == "codex" else "mcpServers" if not isinstance(data, dict) or not isinstance(data.get(key, {}), dict): raise ValueError("invalid object") except (ValueError, UnicodeError, RecursionError): raise SetupError("The client configuration is invalid. Repair it before running setup; no registration was attempted.") from None return data def config_path(client, home): if client == "agy": return home / ".gemini" / "config" / "mcp_config.json" if client == "claude": directory = os.environ.get("CLAUDE_CONFIG_DIR") return Path(directory).expanduser().absolute() / ".claude.json" if directory else home / ".claude.json" return Path(os.environ.get("CODEX_HOME", str(home / ".codex"))).expanduser().absolute() / "config.toml" def powershell_arguments(arguments): """Encode command syntax, never credentials, for Windows cmd/Git Bash callers.""" literal = lambda value: "'" + str(value).replace("'", "''") + "'" script = "$ErrorActionPreference='Stop'; & " + " ".join(map(literal, arguments)) + "; exit $LASTEXITCODE" return [windows_files().system_powershell(), "-NoLogo", "-NoProfile", "-NonInteractive", "-EncodedCommand", base64.b64encode(script.encode("utf-16-le")).decode("ascii")] def helper_command(helper): arguments = [sys.executable, str(helper), "--headers"] if os.name == "nt": command = powershell_arguments(arguments) # Both supported Windows hosts wrap this string for cmd.exe /C. An # absolute quoted system executable avoids project-directory shadowing. return '"' + command[0] + '" ' + " ".join(command[1:]) return shlex.join(arguments) def check_codex_http_version(executable): """Require the version validated with dynamic HTTP headers for this setup.""" try: # The native CLI can create its configuration directory even for --version. arguments = powershell_arguments([executable, "--version"]) if os.name == "nt" else [executable, "--version"] result = subprocess.run(arguments, capture_output=True, timeout=10, check=False, **({"umask": 0o077} if os.name == "posix" else {})) match = re.fullmatch(rb"codex-cli (\d{1,4})\.(\d{1,4})\.(\d{1,4})[ \r\n]*", result.stdout) if result.returncode == 0 and match and tuple(int(part) for part in match.groups()) >= (0, 154, 0): return except (OSError, subprocess.TimeoutExpired): pass raise SetupError("This HTTP setup requires Codex 0.154.0 or newer. Update Codex or use --transport stdio; no registration was attempted.") def entry_for(client, transport, token_path, helper, uv): if transport == "http": command = helper_command(helper) if client == "claude": return {"type": "http", "url": MCP_URL, "headersHelper": command} return {"url": MCP_URL, "tool_timeout_sec": 180, "http_headers_helper": command} entry = { "command": uv, "args": ["tool", "run", "--python", "3.12", "--from", PACKAGE, "vt-mcp"], "env": {"VTAI_TOKEN_FILE": str(token_path), "VTAI_BASE_URL": f"https://{HOST}/api/v3"}, } if client == "codex": entry["tool_timeout_sec"] = 180 else: entry["type"] = "stdio" return entry def updated_config(raw, client, entry): data = parse_config(raw, client) key = "mcp_servers" if client == "codex" else "mcpServers" servers = data.get(key, {}) if "virustotal" in servers: if servers["virustotal"] != entry: raise SetupError("An existing virustotal entry uses different settings. Review and rename or remove only that entry in your client settings, then rerun. Existing settings were preserved.") return raw if client != "codex": data.setdefault(key, {})["virustotal"] = entry return (json.dumps(data, ensure_ascii=True, indent=2) + "\n").encode() # Append only our table, preserving every existing byte and TOML comment. lines = ["", "[mcp_servers.virustotal]"] for name, value in entry.items(): if isinstance(value, dict): continue lines.append(f"{name} = {json.dumps(value, ensure_ascii=False)}") for name, values in entry.items(): if isinstance(values, dict): lines.extend(["", f"[mcp_servers.virustotal.{name}]"]) lines.extend(f"{json.dumps(k)} = {json.dumps(v, ensure_ascii=False)}" for k, v in values.items()) updated = (raw or b"") + ("\n".join(lines) + "\n").encode() expected = dict(data) expected[key] = {**servers, "virustotal": entry} if parse_config(updated, client) != expected: raise SetupError("This TOML layout cannot be extended safely. Add a separate MCP table manually.") return updated def register_token(client): """One finite POST; no redirects, proxy environment, body logging or retries.""" connection = http.client.HTTPSConnection(HOST, timeout=20) try: payload = json.dumps({"agent_family": f"vt-mcp-{client}", "agent_version": PACKAGE_VERSION, "display_name": "VirusTotal MCP"}) connection.request("POST", "/api/v3/agents/register", body=payload, headers={ "Content-Type": "application/json", "Accept": "application/json", "User-Agent": "vt-mcp-setup/1", }) response = connection.getresponse() raw = response.read(MAX_RESPONSE_BYTES + 1) if response.status != 200: if response.status == 429: delay = response.getheader("Retry-After", "") wait = f" Wait at least {delay} seconds." if delay.isascii() and delay.isdigit() and len(delay) <= 8 else "" raise RegistrationLimited("Registration was rate limited; no identity was created and no automatic retry was made." + wait + " Reuse an existing token or rerun setup later.") raise SetupError("Registration did not return a usable token. The outcome may be uncertain; rerunning will not create another identity. Recover an existing token and use --token-file.", "registration_uncertain") if len(raw) > MAX_RESPONSE_BYTES: raise ValueError("oversized response") data = json.loads(raw, object_pairs_hook=unique_object) if not isinstance(data, dict) or not isinstance(data.get("agent_token"), str): raise ValueError("missing token") return valid_token(data["agent_token"].encode("ascii")) except SetupError: raise except (OSError, http.client.HTTPException, ValueError, UnicodeError, RecursionError): raise SetupError("Registration could not be completed safely. Its outcome may be uncertain; rerunning will not create another identity. Recover an existing token and use --token-file.", "registration_uncertain") from None finally: connection.close() def setup(client, transport, import_path=None, *, quiet=False): if os.name not in ("posix", "nt"): raise SetupError("Automatic setup supports Windows, macOS and Linux.") if client == "agy" and transport != "stdio": raise SetupError("Agy setup currently supports stdio. Use --transport stdio.") executable = shutil.which(client) if not executable: raise SetupError(f"Install the {client} CLI first, then rerun setup.") if os.name == "nt" and transport == "http": windows_files().system_powershell() if client == "codex" and transport == "http": check_codex_http_version(executable) if transport == "stdio" and "VTAI_TOKEN" in os.environ: raise SetupError("VTAI_TOKEN conflicts with the protected token file. Unset VTAI_TOKEN before setup and when starting the client; its value was not read.") uv = shutil.which("uv") if transport == "stdio" else None if transport == "stdio" and not uv: raise SetupError("Install uv first (https://docs.astral.sh/uv/getting-started/installation/), then rerun setup.") if os.name == "nt" and uv and Path(uv).suffix.lower() != ".exe": raise SetupError("Windows stdio setup requires the native uv.exe. Install uv's Windows executable; batch wrappers are not supported.") uv = str(Path(uv).absolute()) if uv else None home = Path.home().absolute() if os.name == "nt" else Path.home().resolve() directory = home / ".config" / "vt-mcp" token_path = directory / "token" helper = directory / "setup.py" destination = config_path(client, home) with setup_lock(directory): private_directory(destination.parent) before = read_file(destination) token = token_from_file(token_path) if import_path: imported = token_from_file(Path(import_path).expanduser().absolute()) if imported is None: raise SetupError("The import token file does not exist; no registration was attempted.") if token and imported != token: raise SetupError("A different local token already exists. Existing credentials were preserved; review it before importing another.") token = imported token_before = read_file(token_path, 1024, private=True) entry = entry_for(client, transport, token_path, helper, uv) updated_config(before, client, entry) # Validate all conflicts before registering. helper_before = None helper_bytes = None if transport == "http": helper_bytes = read_file(Path(__file__).absolute(), 256 * 1024) helper_before = read_file(helper, 256 * 1024) if helper_bytes is None or (helper_before is not None and helper_before != helper_bytes and hashlib.sha256(helper_before).hexdigest() not in MANAGED_HELPER_SHA256): raise SetupError("The persistent setup helper conflicts with an existing file. Review ~/.config/vt-mcp/setup.py before continuing.") pending = directory / "registration-pending" if token is None: if read_file(pending, 1024) is not None: raise SetupError("An earlier registration may have succeeded. Setup will not register again. Recover its token and rerun with --token-file; only remove ~/.config/vt-mcp/registration-pending after confirming no identity was created.", "registration_uncertain") atomic_write(pending, b"Registration requested; retain until its token is stored.\n", None) try: token = register_token(client) except RegistrationLimited: pending.unlink() raise if token_before is None: atomic_write(token_path, (token + "\n").encode(), None) # A durable token makes restart safe even if a later configuration write fails. with contextlib.suppress(FileNotFoundError): pending.unlink() if helper_bytes is not None and helper_before != helper_bytes: atomic_write(helper, helper_bytes, helper_before) entry = entry_for(client, transport, token_path, helper, uv) after = updated_config(before, client, entry) if after != before: atomic_write(destination, after, before) elif before is not None and os.name == "posix": # Keep the requested client configuration private on repeat runs too. with os.fdopen(os.open(destination, os.O_RDONLY | os.O_NOFOLLOW), "rb") as config: os.fchmod(config.fileno(), 0o600) if quiet: return print(f"VirusTotal configured for {client} ({transport}). Your token was saved locally; it is not shown.") print(f"Restart {client}, inspect its MCP servers, and ask: Check virustotal.com with VirusTotal.") if client == "claude" and transport == "http": print("Claude Code runs the credential helper in trusted workspaces. Accept its normal trust prompt once when needed.") if transport == "stdio": print("The first connection downloads pinned vt-mcp and Python 3.12 through uv.") def diagnostic(state, client, transport, next_action, error=None): return {"state": state, "client": client, "transport": transport, "next_action": next_action, "error": error} def configured_package(entry): """Recognize a pinned uv command without running it or inspecting its cache.""" command, args = entry.get("command"), entry.get("args") if not isinstance(command, str) or not isinstance(args, list) or len(args) != 7: return None program = Path(command).name if os.name == "nt": program = program.lower().removesuffix(".exe") if program != "uv" or args[:5] != ["tool", "run", "--python", "3.12", "--from"] or args[6] != "vt-mcp": return None match = re.fullmatch(r"vt-mcp==(\d{1,4}\.\d{1,4}\.\d{1,4})", args[5]) if isinstance(args[5], str) else None return match[1] if match else None def update_command(client): python = "py -3" if os.name == "nt" else "python3" return f"{python} setup.py --client {client} --update --json" def managed_helper_command(command, helper): if not isinstance(command, str): return False if os.name == "nt": return command == helper_command(helper) try: parts = shlex.split(command) except ValueError: return False return len(parts) == 3 and parts[1:] == [str(helper), "--headers"] and bool(shutil.which(parts[0])) def replace_package_pin(raw, client, entry, version): """Change only a recognized package argument, validating the whole parsed result.""" data = parse_config(raw, client) key = "mcp_servers" if client == "codex" else "mcpServers" updated = dict(entry, args=list(entry["args"])) updated["args"][5] = PACKAGE data[key]["virustotal"] = updated if client != "codex": return (json.dumps(data, ensure_ascii=True, indent=2) + "\n").encode() # The installer emits one ordinary TOML table. Custom layouts fail closed; # comments and every unrelated byte survive a supported update. table = re.search(rb"(?m)^\[mcp_servers\.virustotal\][ \t]*(?:#[^\r\n]*)?\r?$", raw) if table: tail = raw[table.end():] following = re.search(rb"(?m)^\[", tail) end = table.end() + (following.start() if following else len(tail)) block = raw[table.end():end] pattern = rb"(['\"])" + re.escape(("vt-mcp==" + version).encode()) + rb"\1" matches = list(re.finditer(pattern, block)) if len(matches) == 1: match = matches[0] changed = block[:match.start()] + match[1] + PACKAGE.encode() + match[1] + block[match.end():] candidate = raw[:table.end()] + changed + raw[end:] if parse_config(candidate, client) == data: return candidate raise SetupError("This TOML layout needs a manual package-pin update. Preserve the existing token and other settings.", "custom_configuration") def update_setup(client, requested_transport=None): """Maintain an existing managed entry; never register, import or validate tokens online.""" if os.name not in ("posix", "nt"): raise SetupError("Automatic updates support Windows, macOS and Linux.") home = Path.home().absolute() if os.name == "nt" else Path.home().resolve() destination = config_path(client, home) before = read_file(destination) entry = parse_config(before, client).get("mcp_servers" if client == "codex" else "mcpServers", {}).get("virustotal") if not isinstance(entry, dict): raise SetupError("No existing virustotal entry was found. Use first-time setup only when you need a new connection.", "server_not_configured") transport = "http" if "url" in entry else "stdio" if requested_transport and transport != requested_transport: raise SetupError("Rerun --update without --transport to maintain the configured connection.", "transport_mismatch") if entry.get("enabled") is False or entry.get("disabled") is True: raise SetupError("The virustotal entry is disabled. Review it in your client before updating.", "server_disabled") directory = home / ".config" / "vt-mcp" helper = directory / "setup.py" helper_before = helper_bytes = None after = before if transport == "http": if entry.get("url") != MCP_URL or client == "agy": raise SetupError("This endpoint is custom. Update it in your client; existing settings were preserved.", "custom_configuration") if not any(name in entry for name in AUTH_FIELDS): return diagnostic("remote_managed", client, transport, "The remote server is updated by VirusTotal. Reconnect in your client and inspect its tools; OAuth credentials and tool access were not tested. No settings or credentials were changed.") key = "http_headers_helper" if client == "codex" else "headersHelper" if not managed_helper_command(entry.get(key), helper) or any(name in entry for name in AUTH_FIELDS - {key}): raise SetupError("This authentication setup is custom. Preserve its credential mapping and update it in your client.", "custom_configuration") helper_before = read_file(helper, 256 * 1024) helper_bytes = read_file(Path(__file__).absolute(), 256 * 1024) if helper_bytes is None or helper_before is None or (helper_before != helper_bytes and hashlib.sha256(helper_before).hexdigest() not in MANAGED_HELPER_SHA256): raise SetupError("The local helper is not a recognized installer copy. Preserve it and review the manual setup guide.", "custom_helper") token_path = directory / "token" else: version = configured_package(entry) if version is None: raise SetupError("This server command is custom. Update its package with the original package manager; preserve its configuration and token.", "custom_configuration") if tuple(map(int, version.split("."))) > tuple(map(int, PACKAGE_VERSION.split("."))): raise SetupError("The configured package is newer than this installer. Download the current installer; no downgrade was made.", "newer_package") env = entry.get("env") if not isinstance(env, dict) or env.get("VTAI_BASE_URL") != f"https://{HOST}/api/v3" or "VTAI_TOKEN" in env or "VTAI_TOKEN" in os.environ: raise SetupError("Review the existing token-file and API-origin mapping before updating; do not register again.", "custom_configuration") value = env.get("VTAI_TOKEN_FILE") if not isinstance(value, str) or not value or not Path(value).expanduser().is_absolute(): raise SetupError("Use the existing absolute protected token-file path before updating.", "custom_configuration") token_path = Path(value).expanduser() if version != PACKAGE_VERSION: after = replace_package_pin(before, client, entry, version) if token_from_file(token_path) is None: raise SetupError("Restore the existing protected token before updating. Update never creates another identity.", "missing_token") if not shutil.which(client): raise SetupError(f"Install the {client} CLI before updating its existing configuration.", "client_not_found") if transport == "stdio" and (not shutil.which(entry["command"]) or os.name == "nt" and Path(entry["command"]).suffix.lower() != ".exe"): raise SetupError("Restore the configured native uv executable before updating.", "server_command_not_found") changed = after != before or helper_bytes is not None and helper_bytes != helper_before with setup_lock(directory): if read_file(destination) != before or token_from_file(token_path) is None: raise SetupError("Local configuration changed during update. Recheck it before trying again.", "concurrent_change") if helper_bytes is not None and helper_bytes != helper_before: atomic_write(helper, helper_bytes, helper_before) if after != before: atomic_write(destination, after, before) return diagnostic("updated" if changed else "up_to_date", client, transport, "Existing credentials and other settings were preserved. Restart the client, run --check, then make a VirusTotal tool call. " + ("uv downloads the configured package on the next connection; installation and runtime were not tested." if transport == "stdio" else "The remote server updates independently; this only maintains the local credential helper.")) def credential_status(token): """Check authentication only, without a lookup, quota admission or retry.""" connection = http.client.HTTPSConnection(HOST, timeout=20) try: connection.request("GET", "/api/v3/agents/me/access", headers={ "Authorization": "Bearer " + token, "Accept": "application/json", "User-Agent": "vt-mcp-setup/1", }) response = connection.getresponse() raw = response.read(MAX_RESPONSE_BYTES + 1) if response.status in (401, 403): return "auth_required", "credential_rejected" if response.status == 429: return "service_unavailable", "rate_limited" if response.status != 200: return "service_unavailable", "access_check_unavailable" if len(raw) > MAX_RESPONSE_BYTES: return "service_unavailable", "invalid_access_response" data = json.loads(raw, object_pairs_hook=unique_object) if data != {"status": "active", "authentication": "agent_token"}: return "service_unavailable", "invalid_access_response" return "credential_active", None except (OSError, http.client.HTTPException, ValueError, UnicodeError, RecursionError): return "service_unavailable", "access_check_unavailable" finally: connection.close() def check_setup(client, requested_transport=None): """Read user settings only; never run clients/helpers or modify local files.""" transport = requested_transport or ("stdio" if client == "agy" else "http") maintenance = {"management": "unknown", "configured_package_version": None, "target_package_version": None, "installed_package_version": None, "runtime_verified": False, "tool_call_verified": False, "update_available": None, "update_command": None, "setup_url": f"https://{HOST}/setup.py"} def result(state, action, error=None): return dict(diagnostic(state, client, transport, action, error), maintenance=dict(maintenance)) if os.name not in ("posix", "nt"): return result("configuration_unverified", "Inspect MCP status in your client; automatic checks support Windows, macOS and Linux.", "unsupported_platform") if not shutil.which(client): return result("missing_client", f"Install the {client} CLI, then rerun this check.", "client_not_found") home = Path.home().absolute() if os.name == "nt" else Path.home().resolve() try: raw = read_file(config_path(client, home)) data = parse_config(raw, client) except (SetupError, OSError): return result("configuration_error", "Repair the client user configuration and its file permissions, then rerun this check.", "invalid_configuration") entry = data.get("mcp_servers" if client == "codex" else "mcpServers", {}).get("virustotal") if entry is None: return result("not_configured", "Configure VirusTotal in this client's user settings, or inspect its project-scoped MCP configuration.", "server_not_configured") if not isinstance(entry, dict): return result("configuration_error", "Repair the virustotal entry in your client settings.", "invalid_server_entry") if entry.get("enabled") is False or entry.get("disabled") is True: return result("configuration_error", "Enable the virustotal MCP entry in your client settings.", "server_disabled") actual_transport = "http" if "url" in entry else "stdio" if "command" in entry else None if actual_transport is None: return result("configuration_unverified", "Inspect this custom MCP configuration in the native client.", "unsupported_configuration") if requested_transport and actual_transport != requested_transport: return result("configuration_error", "Rerun --check without --transport to inspect the configured transport.", "transport_mismatch") transport = actual_transport directory = home / ".config" / "vt-mcp" token_path = directory / "token" if transport == "http": if entry.get("url") != MCP_URL or client == "agy": return result("configuration_error", "Use the documented VirusTotal endpoint and a supported client transport.", "unsupported_endpoint") helper_key = "http_headers_helper" if client == "codex" else "headersHelper" auth_fields = AUTH_FIELDS if not any(name in entry for name in auth_fields): maintenance.update(management="remote_oauth", update_available=False) action = "Open /mcp in Claude Code to sign in and check the connection." if client == "claude" else "Run codex mcp login virustotal, then inspect MCP status in Codex." return result("authentication_unverified", action + " The remote server updates independently. Native OAuth credentials and client trust were not inspected.", "native_oauth_not_checked") try: configured_helper = entry.get(helper_key, "") if not isinstance(configured_helper, str): return result("configuration_error", "Repair the credential helper command in your client settings.", "invalid_helper") if os.name == "nt": managed = configured_helper == helper_command(directory / "setup.py") command = [windows_files().system_powershell()] else: command = shlex.split(configured_helper) managed = len(command) == 3 and command[1:] == [str(directory / "setup.py"), "--headers"] if not managed or any(name in entry for name in auth_fields - {helper_key}): return result("configuration_unverified", "Inspect authentication in the native client; this check supports the protected-file helper installed by setup.", "custom_authentication") interpreter_present = Path(command[0]).is_file() if os.name == "nt" else shutil.which(command[0]) helper_bytes = read_file(directory / "setup.py", 256 * 1024) if not interpreter_present or not helper_bytes: return result("configuration_error", "Restore the configured Python interpreter and local credential helper before reconnecting.", "missing_helper") current_bytes = read_file(Path(__file__).absolute(), 256 * 1024) known = helper_bytes == current_bytes or hashlib.sha256(helper_bytes).hexdigest() in MANAGED_HELPER_SHA256 maintenance.update(management="managed_http" if known else "custom_helper", update_available=helper_bytes != current_bytes if known else None, update_command=update_command(client) if known else None) except (SetupError, OSError, ValueError, TypeError): return result("configuration_error", "Repair the credential helper configuration and file permissions.", "invalid_helper") else: env = entry.get("env", {}) if not isinstance(env, dict) or not isinstance(entry.get("command"), str): return result("configuration_error", "Repair the command and environment in the virustotal entry.", "invalid_server_entry") base = env.get("VTAI_BASE_URL", os.environ.get("VTAI_BASE_URL", f"https://{HOST}/api/v3")) if base != f"https://{HOST}/api/v3": return result("configuration_error", "Set VTAI_BASE_URL to the documented VirusTotal API origin in the client entry.", "unsupported_endpoint") if "VTAI_TOKEN" in env or "VTAI_TOKEN" in os.environ: return result("auth_required", "Unset VTAI_TOKEN when using VTAI_TOKEN_FILE, then rerun this check.", "conflicting_credentials") if not shutil.which(entry["command"]): return result("configuration_error", "Install the executable named in your MCP configuration, then rerun this check.", "server_command_not_found") program = Path(entry["command"]).name if os.name == "nt": program = program.lower().removesuffix(".exe") arguments = entry.get("args", []) version = configured_package(entry) if not (program == "vt-mcp" and arguments == [] or version is not None): return result("configuration_unverified", "Inspect this custom server command in the native client; no program was executed by this check.", "custom_server_command") maintenance.update(management="managed_stdio" if version else "manual_package", configured_package_version=version, target_package_version=PACKAGE_VERSION, update_available=tuple(map(int, version.split("."))) < tuple(map(int, PACKAGE_VERSION.split("."))) if version else None, update_command=update_command(client) if version else None) value = env.get("VTAI_TOKEN_FILE") if not isinstance(value, str) or not value: return result("configuration_unverified", "Inspect this custom credential configuration in the native client.", "custom_authentication") token_path = Path(value).expanduser() if not token_path.is_absolute(): return result("configuration_unverified", "Use an absolute VTAI_TOKEN_FILE path so it is independent of the client's working directory.", "relative_token_path") try: token = token_from_file(token_path) except (SetupError, OSError, ValueError): return result("auth_required", "Check the token file's format, ownership and owner-only permissions; do not register again.", "invalid_token_file") if token is None: return result("auth_required", "Restore the existing protected token file; do not create another identity just to reconnect.", "missing_token") state, error = credential_status(token) if state == "credential_active": action = "User configuration and Agent Token checked. Package versions describe configuration only; installed package, runtime, tools, client trust, project overrides, quota and upstream availability were not tested." if maintenance["update_available"]: action += f" Download https://{HOST}/setup.py, then run: {update_command(client)}." elif maintenance["management"] in {"manual_package", "custom_helper"}: action += " Maintain this custom installation with its original setup instructions, preserving the token; then make a tool call in the client." else: action += " Run a first VirusTotal tool call in the native client." return result(state, action) if state == "auth_required": return result(state, "The Agent Token was rejected. Restore valid access or use native OAuth; do not retry registration automatically.", error) return result(state, "Retry this read-only check later. Existing credentials and configuration were preserved; no registration or lookup was attempted.", error) def check_exit_code(result): if result["state"] == "credential_active": return 0 if result["state"] in {"authentication_unverified", "configuration_unverified"}: return 2 return 1 def print_diagnostic(result, json_output): if json_output: print(json.dumps(result, ensure_ascii=True, separators=(",", ":"))) else: print(f"VirusTotal {result['state']}: {result['next_action']}") def emit_headers(): """The client invokes this mode; its stdout is the credential channel.""" token = token_from_file(Path.home() / ".config" / "vt-mcp" / "token") if token is None: raise SetupError("The local VTAI token is missing. Import it with setup --token-file.") print(json.dumps({"Authorization": "Bearer " + token})) def main(argv=None): arguments = list(sys.argv[1:] if argv is None else argv) json_output = "--json" in arguments parser = argparse.ArgumentParser(description=__doc__) modes = parser.add_mutually_exclusive_group(required=True) modes.add_argument("--client", choices=("agy", "claude", "codex")) modes.add_argument("--headers", action="store_true", help=argparse.SUPPRESS) parser.add_argument("--transport", choices=("http", "stdio")) parser.add_argument("--token-file", type=Path, help="import an existing owner-only token file instead of registering") actions = parser.add_mutually_exclusive_group() actions.add_argument("--check", action="store_true", help="check user configuration and Agent Token without changing files or making a lookup") actions.add_argument("--update", action="store_true", help="update a recognized existing configuration, preserving credentials; never register") parser.add_argument("--json", action="store_true", help="print one safe JSON result (check exits: 0 active token, 1 action needed, 2 native verification needed)") try: with contextlib.redirect_stderr(io.StringIO()) if json_output else contextlib.nullcontext(): args = parser.parse_args(arguments) if args.headers and (args.transport or args.token_file or args.check or args.update or args.json): parser.error("--headers cannot be combined with setup or diagnostic options") if (args.check or args.update) and args.token_file: parser.error("--check and --update cannot import a token; run setup separately") except SystemExit as error: if not json_output or error.code == 0: raise print_diagnostic(diagnostic("invalid_arguments", None, None, "Run --help for supported setup and diagnostic options.", "invalid_arguments"), True) return 1 transport = args.transport or (None if args.update else "stdio" if args.client == "agy" else "http") try: if args.headers: emit_headers() elif args.check: result = check_setup(args.client, args.transport) print_diagnostic(result, args.json) return check_exit_code(result) elif args.update: print_diagnostic(update_setup(args.client, args.transport), args.json) else: setup(args.client, transport, args.token_file, quiet=args.json) if args.json: print_diagnostic(diagnostic("configured", args.client, transport, "Run --check to validate the Agent Token, then make a first VirusTotal tool call in your client. Configuration alone does not verify a connection."), True) except (SetupError, OSError) as error: message = str(error) if isinstance(error, SetupError) else "A local setup file could not be accessed. Check permissions and rerun; existing credentials will be reused." if json_output: code = "rate_limited" if isinstance(error, RegistrationLimited) else error.code if isinstance(error, SetupError) else "local_file_error" print_diagnostic(diagnostic("setup_failed", args.client, transport, message, code), True) else: print(f"VirusTotal setup: {message}", file=sys.stderr) return 1 except KeyboardInterrupt: message = "Operation interrupted. Rerun safely; an uncertain registration will not be repeated." if json_output: print_diagnostic(diagnostic("interrupted", args.client, transport, message, "interrupted"), True) else: print("VirusTotal setup: " + message, file=sys.stderr) return 1 return 0 if __name__ == "__main__": raise SystemExit(main())