Merge pull request #26 from GDATASoftwareAG/add-oauth-in-python

add connect_with_client_credentials

Author: doxthree

.gitignore

@@ -39,6 +39,7 @@ java/.project
 java/.settings
 java/bin
 lib/
+*.crt
 
 bin/
 obj/

python/requirements.txt

@@ -4,4 +4,5 @@ websockets~=10.3
 python-dotenv==0.20.0
 httpx[http2]==0.22.0
 build==0.7.0
-jwt==1.3.1
+jwt==1.3.1
+authlib==1.0.1

python/setup.cfg

@@ -20,6 +20,7 @@ install_requires =
     websockets == 10.3
     httpx[http2] == 0.22.0
     jwt == 1.3.1
+    authlib == 1.0.1
 
 [options.packages.find]
 where = src

python/src/vaas/vaas.py

@@ -10,9 +10,12 @@
 from jwt import JWT
 import httpx
 import websockets.client
+from authlib.integrations.httpx_client import AsyncOAuth2Client
+
 
 URL = "wss://gateway-vaas.gdatasecurity.de"
 
+
 class VaasTracing:
     """Tracing interface for Vaas"""
 
@@ -37,7 +40,7 @@ def __init__(self, tracing=VaasTracing()):
     async def connect(self, token, url=URL):
         """Connect to VaaS
 
-        token -- a OpenID Connect token signed by a trusted identity provider
+        token -- OpenID Connect token signed by a trusted identity provider
         """
         self.websocket = await websockets.client.connect(url)
         authenticate_request = {"kind": "AuthRequest", "token": token}
@@ -53,6 +56,22 @@ async def connect(self, token, url=URL):
             self.__receive_loop()
         )  # fire and forget async_foo()
 
+    async def connect_with_client_credentials(
+        self, client_id, client_secret, token_endpoint, url=URL, verify=True
+    ):
+        """Connect to VaaS with client credentials grant
+
+        :param str client_id: Client ID provided by G DATA
+        :param str client_secret: Client secret provided by G DATA
+        :param str token_endpoint: Token endpoint of identity provider
+        :param str url: Websocket endpoint for verdict requests
+        :param bool verify: This switch turns off SSL validation when set to False; default: True
+
+        """
+        async with AsyncOAuth2Client(client_id, client_secret, verify=verify) as client:
+            token = (await client.fetch_token(token_endpoint))["access_token"]
+        await self.connect(token, url)
+
     async def close(self):
         """Close the connection"""
         if self.websocket is not None:

python/test.py

@@ -1,7 +1,7 @@
 import unittest
 import xmlrunner
 
-from tests.test_vaas import VaasTest # pylint: disable=unused-import
+from tests.test_vaas import VaasTest  # pylint: disable=unused-import
 
 if __name__ == "__main__":
     unittest.main(testRunner=xmlrunner.XMLTestRunner(output="test-reports"))

python/tests/test_vaas.py

@@ -73,5 +73,6 @@ async def test_for_buffer_traces(self):
             tracing.trace_hash_request.assert_called_with(ANY)
             tracing.trace_upload_request.assert_called_with(ANY, 1024)
 
+
 if __name__ == "__main__":
     unittest.main()