Migrate to using aiohttp

Author: bdraco

onvif/client.py

@@ -6,20 +6,22 @@
 import datetime as dt
 import logging
 import os.path
-from typing import Any
 from collections.abc import Callable
-import httpx
-from httpx import AsyncClient, BasicAuth, DigestAuth
+from typing import Any
+
+import zeep.helpers
 from zeep.cache import SqliteCache
 from zeep.client import AsyncClient as BaseZeepAsyncClient
-import zeep.helpers
 from zeep.proxy import AsyncServiceProxy
-from zeep.transports import AsyncTransport
 from zeep.wsdl import Document
 from zeep.wsse.username import UsernameToken
 
+import aiohttp
+import httpx
+from aiohttp import BasicAuth, ClientSession, DigestAuthMiddleware, TCPConnector
 from onvif.definition import SERVICES
 from onvif.exceptions import ONVIFAuthError, ONVIFError, ONVIFTimeoutError
+from requests import Response
 
 from .const import KEEPALIVE_EXPIRY
 from .managers import NotificationManager, PullPointManager
@@ -29,13 +31,14 @@
 from .util import (
     create_no_verify_ssl_context,
     normalize_url,
+    obscure_user_pass_url,
     path_isfile,
-    utcnow,
     strip_user_pass_url,
-    obscure_user_pass_url,
+    utcnow,
 )
-from .wrappers import retry_connection_error  # noqa: F401
+from .wrappers import retry_connection_error
 from .wsa import WsAddressingIfMissingPlugin
+from .zeep_aiohttp import AIOHTTPTransport
 
 logger = logging.getLogger("onvif")
 logging.basicConfig(level=logging.INFO)
@@ -48,7 +51,7 @@
 _CONNECT_TIMEOUT = 30
 _READ_TIMEOUT = 90
 _WRITE_TIMEOUT = 90
-_HTTPX_LIMITS = httpx.Limits(keepalive_expiry=KEEPALIVE_EXPIRY)
+_KEEPALIVE_TIMEOUT = aiohttp.ClientTimeout(sock_keepalive=KEEPALIVE_EXPIRY)
 _NO_VERIFY_SSL_CONTEXT = create_no_verify_ssl_context()
 
 
@@ -59,7 +62,7 @@ def wrapped(*args, **kwargs):
         try:
             return func(*args, **kwargs)
         except Exception as err:
-            raise ONVIFError(err)
+            raise ONVIFError(err) from err
 
     return wrapped
 
@@ -80,7 +83,7 @@ def __init__(self, user, passw, dt_diff=None, **kwargs):
     def apply(self, envelope, headers):
         old_created = self.created
         if self.created is None:
-            self.created = dt.datetime.now(tz=dt.timezone.utc).replace(tzinfo=None)
+            self.created = dt.datetime.now(tz=dt.UTC).replace(tzinfo=None)
         if self.dt_diff is not None:
             self.created += self.dt_diff
         result = super().apply(envelope, headers)
@@ -102,20 +105,28 @@ def original_load(self, *args: Any, **kwargs: Any) -> None:
         return original_load(self, *args, **kwargs)
 
 
-class AsyncTransportProtocolErrorHandler(AsyncTransport):
-    """Retry on remote protocol error.
+class AsyncTransportProtocolErrorHandler(AIOHTTPTransport):
+    """
+    Retry on remote protocol error.
 
     http://datatracker.ietf.org/doc/html/rfc2616#section-8.1.4 allows the server
     # to close the connection at any time, we treat this as a normal and try again
     # once since
     """
 
-    @retry_connection_error(attempts=2, exception=httpx.RemoteProtocolError)
-    async def post(self, address, message, headers):
+    @retry_connection_error(attempts=2, exception=aiohttp.ServerDisconnectedError)
+    async def post(
+        self, address: str, message: str, headers: dict[str, str]
+    ) -> httpx.Response:
         return await super().post(address, message, headers)
 
-    @retry_connection_error(attempts=2, exception=httpx.RemoteProtocolError)
-    async def get(self, address, params, headers):
+    @retry_connection_error(attempts=2, exception=aiohttp.ServerDisconnectedError)
+    async def get(
+        self,
+        address: str,
+        params: dict[str, Any] | None = None,
+        headers: dict[str, str] | None = None,
+    ) -> Response:
         return await super().get(address, params, headers)
 
 
@@ -162,17 +173,18 @@ def __init__(self, *args, **kwargs):
         self.set_ns_prefix("wsa", "http://www.w3.org/2005/08/addressing")
 
     def create_service(self, binding_name, address):
-        """Create a new ServiceProxy for the given binding name and address.
+        """
+        Create a new ServiceProxy for the given binding name and address.
         :param binding_name: The QName of the binding
         :param address: The address of the endpoint
         """
         try:
             binding = self.wsdl.bindings[binding_name]
         except KeyError:
             raise ValueError(
-                "No binding found with the given QName. Available bindings "
-                "are: %s" % (", ".join(self.wsdl.bindings.keys()))
-            )
+                f"No binding found with the given QName. Available bindings "
+                f"are: {', '.join(self.wsdl.bindings.keys())}"
+            ) from None
         return AsyncServiceProxy(self, binding, address=address)
 
 
@@ -223,7 +235,7 @@ def __init__(
         write_timeout: int | None = None,
     ) -> None:
         if not path_isfile(url):
-            raise ONVIFError("%s doesn`t exist!" % url)
+            raise ONVIFError(f"{url} doesn`t exist!")
 
         self.url = url
         self.xaddr = xaddr
@@ -236,28 +248,28 @@ def __init__(
         self.dt_diff = dt_diff
         self.binding_name = binding_name
         # Create soap client
-        timeouts = httpx.Timeout(
-            _DEFAULT_TIMEOUT,
-            connect=_CONNECT_TIMEOUT,
-            read=read_timeout or _READ_TIMEOUT,
-            write=write_timeout or _WRITE_TIMEOUT,
+        timeout_seconds = _DEFAULT_TIMEOUT
+        operation_timeout = read_timeout or _READ_TIMEOUT
+        connector = TCPConnector(
+            ssl=_NO_VERIFY_SSL_CONTEXT,
+            keepalive_timeout=KEEPALIVE_EXPIRY,
         )
-        client = AsyncClient(
-            verify=_NO_VERIFY_SSL_CONTEXT, timeout=timeouts, limits=_HTTPX_LIMITS
+        session = ClientSession(
+            connector=connector,
+            timeout=aiohttp.ClientTimeout(
+                total=timeout_seconds,
+                connect=_CONNECT_TIMEOUT,
+                sock_read=operation_timeout,
+            ),
         )
-        # The wsdl client should never actually be used, but it is required
-        # to avoid creating another ssl context since the underlying code
-        # will try to create a new one if it doesn't exist.
-        wsdl_client = httpx.Client(
-            verify=_NO_VERIFY_SSL_CONTEXT, timeout=timeouts, limits=_HTTPX_LIMITS
-        )
-        self.transport = (
-            AsyncTransportProtocolErrorHandler(client=client, wsdl_client=wsdl_client)
-            if no_cache
-            else AsyncTransportProtocolErrorHandler(
-                client=client, wsdl_client=wsdl_client, cache=SqliteCache()
-            )
+        self.transport = AsyncTransportProtocolErrorHandler(
+            session=session,
+            timeout=timeout_seconds,
+            operation_timeout=operation_timeout,
+            verify_ssl=False,
         )
+        if not no_cache:
+            self.transport.cache = SqliteCache()
         self.document: Document | None = None
         self.zeep_client_authless: ZeepAsyncClient | None = None
         self.ws_client_authless: AsyncServiceProxy | None = None
@@ -399,7 +411,8 @@ def __init__(
         self.to_dict = ONVIFService.to_dict
 
         self._snapshot_uris = {}
-        self._snapshot_client = AsyncClient(verify=_NO_VERIFY_SSL_CONTEXT)
+        self._snapshot_connector = TCPConnector(ssl=_NO_VERIFY_SSL_CONTEXT)
+        self._snapshot_client = ClientSession(connector=self._snapshot_connector)
 
     async def get_capabilities(self) -> dict[str, Any]:
         """Get device capabilities."""
@@ -531,7 +544,7 @@ async def create_notification_manager(
 
     async def close(self) -> None:
         """Close all transports."""
-        await self._snapshot_client.aclose()
+        await self._snapshot_client.close()
         for service in self.services.values():
             await service.close()
 
@@ -572,42 +585,54 @@ async def get_snapshot(
         if uri is None:
             return None
 
+        # Create a new session with appropriate auth
+        connector = TCPConnector(ssl=_NO_VERIFY_SSL_CONTEXT)
+        middlewares = []
         auth = None
+
         if self.user and self.passwd:
             if basic_auth:
                 auth = BasicAuth(self.user, self.passwd)
             else:
-                auth = DigestAuth(self.user, self.passwd)
-
-        response = await self._try_snapshot_uri(uri, auth)
-
-        # If the request fails with a 401, make sure to strip any
-        # sample user/pass from the URL and try again
-        if (
-            response.status_code == 401
-            and (stripped_uri := strip_user_pass_url(uri))
-            and stripped_uri != uri
-        ):
-            response = await self._try_snapshot_uri(stripped_uri, auth)
+                # Use DigestAuthMiddleware for digest auth
+                middlewares.append(DigestAuthMiddleware(self.user, self.passwd))
+
+        async with ClientSession(
+            connector=connector, auth=auth, middlewares=middlewares
+        ) as session:
+            response = await self._try_snapshot_uri_with_session(session, uri)
+            content = await response.read()
+
+            # If the request fails with a 401, make sure to strip any
+            # sample user/pass from the URL and try again
+            if (
+                response.status == 401
+                and (stripped_uri := strip_user_pass_url(uri))
+                and stripped_uri != uri
+            ):
+                response = await self._try_snapshot_uri_with_session(
+                    session, stripped_uri
+                )
+                content = await response.read()
 
-        if response.status_code == 401:
-            raise ONVIFAuthError(f"Failed to authenticate to {uri}")
+            if response.status == 401:
+                raise ONVIFAuthError(f"Failed to authenticate to {uri}")
 
-        if response.status_code < 300:
-            return response.content
+            if response.status < 300:
+                return content
 
-        return None
+            return None
 
-    async def _try_snapshot_uri(
-        self, uri: str, auth: BasicAuth | DigestAuth | None
-    ) -> httpx.Response:
+    async def _try_snapshot_uri_with_session(
+        self, session: ClientSession, uri: str
+    ) -> aiohttp.ClientResponse:
         try:
-            return await self._snapshot_client.get(uri, auth=auth)
-        except httpx.TimeoutException as error:
+            return await session.get(uri)
+        except TimeoutError as error:
             raise ONVIFTimeoutError(
                 f"Timed out fetching {obscure_user_pass_url(uri)}: {error}"
             ) from error
-        except httpx.RequestError as error:
+        except aiohttp.ClientError as error:
             raise ONVIFError(
                 f"Error fetching {obscure_user_pass_url(uri)}: {error}"
             ) from error
@@ -618,7 +643,7 @@ def get_definition(
         """Returns xaddr and wsdl of specified service"""
         # Check if the service is supported
         if name not in SERVICES:
-            raise ONVIFError("Unknown service %s" % name)
+            raise ONVIFError(f"Unknown service {name}")
         wsdl_file = SERVICES[name]["wsdl"]
         namespace = SERVICES[name]["ns"]
 
@@ -629,14 +654,14 @@ def get_definition(
 
         wsdlpath = os.path.join(self.wsdl_dir, wsdl_file)
         if not path_isfile(wsdlpath):
-            raise ONVIFError("No such file: %s" % wsdlpath)
+            raise ONVIFError(f"No such file: {wsdlpath}")
 
         # XAddr for devicemgmt is fixed:
         if name == "devicemgmt":
             xaddr = "{}:{}/onvif/device_service".format(
                 self.host
                 if (self.host.startswith("http://") or self.host.startswith("https://"))
-                else "http://%s" % self.host,
+                else f"http://{self.host}",
                 self.port,
             )
             return xaddr, wsdlpath, binding_name

onvif/zeep_aiohttp.py

@@ -98,24 +98,13 @@ def _aiohttp_to_httpx_response(
         # Store cookies if any
         if aiohttp_response.cookies:
             for cookie in aiohttp_response.cookies.values():
-                # Extract all cookie attributes
-                cookie_attrs = {}
-                if cookie.get("domain"):
-                    cookie_attrs["domain"] = cookie.get("domain")
-                if cookie.get("path"):
-                    cookie_attrs["path"] = cookie.get("path")
-                if cookie.get("secure"):
-                    cookie_attrs["secure"] = True
-                if cookie.get("httponly"):
-                    cookie_attrs["httpOnly"] = True
-                if cookie.get("max-age"):
-                    cookie_attrs["max_age"] = int(cookie.get("max-age"))
-                if cookie.get("expires"):
-                    cookie_attrs["expires"] = cookie.get("expires")
-                if cookie.get("samesite"):
-                    cookie_attrs["samesite"] = cookie.get("samesite")
-
-                httpx_response.cookies.set(cookie.key, cookie.value, **cookie_attrs)
+                # httpx.Cookies.set only accepts name, value, domain, and path
+                httpx_response.cookies.set(
+                    cookie.key,
+                    cookie.value,
+                    domain=cookie.get("domain", ""),
+                    path=cookie.get("path", "/"),
+                )
 
         return httpx_response