Skip to content

Commit

Permalink
feat: Add custom connectors to the client (#173)
Browse files Browse the repository at this point in the history
  • Loading branch information
mgmacias95 authored Dec 19, 2023
1 parent 6bf4dec commit 87ab687
Showing 1 changed file with 8 additions and 1 deletion.
9 changes: 8 additions & 1 deletion vt/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,7 @@ class Client:
made by the client (None by default).
:param headers: Dict of headers defined by the user.
:param verify_ssl: Whether to verify the certificate in SSL connections.
:param connector: (Optional) A custom aiohttp connector.
:type apikey: str
:type agent: str
:type host: str
Expand All @@ -200,6 +201,7 @@ class Client:
:type proxy: str
:type headers: dict
:type verify_ssl: bool
:type connector: aiohttp.BaseConnector
"""

def __init__(
Expand All @@ -212,6 +214,7 @@ def __init__(
proxy: typing.Optional[str] = None,
headers: typing.Optional[typing.Dict] = None,
verify_ssl: bool = True,
connector: aiohttp.BaseConnector = None
):
"""Initialize the client with the provided API key."""

Expand All @@ -230,6 +233,10 @@ def __init__(
self._proxy = proxy
self._user_headers = headers
self._verify_ssl = verify_ssl
if connector is not None:
self._connector = connector
else:
self._connector = aiohttp.TCPConnector(ssl=self._verify_ssl)

def _full_url(self, path:str, *args: typing.Any) -> str:
try:
Expand All @@ -256,7 +263,7 @@ def _get_session(self) -> aiohttp.ClientSession:
headers.update(self._user_headers)

self._session = aiohttp.ClientSession(
connector=aiohttp.TCPConnector(ssl=self._verify_ssl),
connector=self._connector,
headers=headers,
trust_env=self._trust_env,
timeout=aiohttp.ClientTimeout(total=self._timeout),
Expand Down

0 comments on commit 87ab687

Please sign in to comment.