From 8b7a39501b55318b6536ea6d5488a040f06f899a Mon Sep 17 00:00:00 2001
From: Gabriel Cuba <gcuba@whitestack.com>
Date: Wed, 2 Nov 2022 17:21:50 -0500
Subject: [PATCH] Add TLS to gRPC server

---
 osm_ee/frontend_server.py |  3 ++-
 osm_ee/util/util_grpc.py  | 26 ++++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 1 deletion(-)
 create mode 100644 osm_ee/util/util_grpc.py

diff --git a/osm_ee/frontend_server.py b/osm_ee/frontend_server.py
index 079503d..32b3afe 100644
--- a/osm_ee/frontend_server.py
+++ b/osm_ee/frontend_server.py
@@ -33,6 +33,7 @@ from osm_ee.frontend_pb2 import SshKeyRequest, SshKeyReply
 
 from osm_ee.base_ee import BaseEE
 import osm_ee.util.util_ee as util_ee
+import osm_ee.util.util_grpc as util_grpc
 
 
 class FrontendExecutor(FrontendExecutorBase):
@@ -75,7 +76,7 @@ async def main(*, host: str = '0.0.0.0', port: int = 50051) -> None:
     # Start server
     server = Server([FrontendExecutor()])
     with graceful_exit([server]):
-        await server.start(host, port)
+        await server.start(host, port, ssl=util_grpc.create_secure_context())
         logging.getLogger('osm_ee.frontend_server').debug(f'Serving on {host}:{port}')
         await server.wait_closed()
 
diff --git a/osm_ee/util/util_grpc.py b/osm_ee/util/util_grpc.py
new file mode 100644
index 0000000..df904c1
--- /dev/null
+++ b/osm_ee/util/util_grpc.py
@@ -0,0 +1,26 @@
+import logging
+import ssl
+
+logger = logging.getLogger("osm_ee.util_grpc")
+
+SERVER_CERT = "/etc/ssl/grpc-tls/tls.crt"
+SERVER_KEY = "/etc/ssl/grpc-tls/tls.key"
+
+
+def create_secure_context() -> ssl.SSLContext:
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
+    # ctx.verify_mode = ssl.CERT_REQUIRED
+    try:
+        ctx.load_cert_chain(str(SERVER_CERT), str(SERVER_KEY))
+    except FileNotFoundError:
+        logger.warning("TLS Certificate not found, starting gRPC server in unsecure mode")
+        return None
+    # TODO: client TLS 
+    # ctx.load_verify_locations(str(trusted))
+    ctx.set_ciphers('ECDHE+AESGCM:ECDHE+CHACHA20:DHE+AESGCM:DHE+CHACHA20')
+    ctx.set_alpn_protocols(['h2'])
+    try:
+        ctx.set_npn_protocols(['h2'])
+    except NotImplementedError:
+        pass
+    return ctx
\ No newline at end of file
-- 
2.25.1