import vt
import io
import os
import sys
from typing import Optional, Any


class VirusTotal:
	def __init__(self):
		self.API_KEY = os.getenv('VT_API_KEY')
		self.client = vt.Client(self.API_KEY)

	def __enter__(self):
		return self

	def __exit__(self, exc_type, exc_val, exc_tb):
		self.close()

	def close(self):
		if self.client:
			self.client.close()

	def get_content(self, file_hash: str, output_path: Optional[str] = None) -> Optional[bytes]:
		"""
		Downloads file content from VT to either a bytes stream or a file.
		"""
		if output_path is None:
			file_stream = io.BytesIO()
			self.client.download_file(file_hash, file_stream)
			return file_stream.getvalue()
		else:
			with open(output_path, 'wb') as f:
				self.client.download_file(file_hash, f)

	def query_by_hash(self, file_hash: str) -> Any:
		"""
		Returns a VirusTotal object by standard hash (sha256/sha1/md5)
		"""
		return self.client.get_object(f"/files/{file_hash}")

	def query_by_collection(self, collection_id: str) -> list:
		"""
		Returns a list of VirusTotal files object by collection id.
		"""
		try:
			return list(self.client.iterator(f"/collections/{collection_id}/files"))
		except vt.error.APIError as E:
			return []

	def query_custom(self, query: str) -> list:
		"""
		Generic VirusTotal query method
		"""
		return list(self.client.iterator('/intelligence/search', params={'query': query}))

	def query_by_pesha256(self, file_pesha256: str) -> list:
		"""
		Returns a VirusTotal object by authentihash
		"""
		return self.query_custom(f"authentihash:{file_pesha256}")

	def resolve_hash(self, input_hash: str) -> Optional[str]:
		"""
		Attempts to resolve as a standard hash (sha256/sha1/md5), falls back to authentihash
		"""
		try:
			obj = self.query_by_hash(input_hash)
			return obj.id 
		except vt.error.APIError as E:
			if E.code == 'NotFoundError':
				results = self.query_by_pesha256(input_hash)
				if results:
					return results[0].id 
				return None
			else:
				raise E
