From 9b14c2ec9d29c9c5b3c89cbebadada4e7246e3da Mon Sep 17 00:00:00 2001
From: PT-ATA No One <ata-no-one@gdata.de>
Date: Tue, 24 May 2022 11:11:34 +0200
Subject: [PATCH] add connect_with_client_credentials

---
 .gitignore                |  1 +
 python/requirements.txt   |  3 ++-
 python/setup.cfg          |  1 +
 python/src/vaas/vaas.py   | 21 ++++++++++++++++++++-
 python/test.py            |  2 +-
 python/tests/test_vaas.py |  1 +
 6 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 1f861f3e..1ce4210f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -39,6 +39,7 @@ java/.project
 java/.settings
 java/bin
 lib/
+*.crt
 
 bin/
 obj/
diff --git a/python/requirements.txt b/python/requirements.txt
index 1cb20f02..885d2dcc 100644
--- a/python/requirements.txt
+++ b/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
\ No newline at end of file
+jwt==1.3.1
+authlib==1.0.1
\ No newline at end of file
diff --git a/python/setup.cfg b/python/setup.cfg
index 067b4d73..82a66ad4 100644
--- a/python/setup.cfg
+++ b/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
\ No newline at end of file
diff --git a/python/src/vaas/vaas.py b/python/src/vaas/vaas.py
index 68d66557..abbe95e8 100644
--- a/python/src/vaas/vaas.py
+++ b/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:
diff --git a/python/test.py b/python/test.py
index 693dfee5..aef8820e 100644
--- a/python/test.py
+++ b/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"))
diff --git a/python/tests/test_vaas.py b/python/tests/test_vaas.py
index 5e20470b..32800f9c 100644
--- a/python/tests/test_vaas.py
+++ b/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()