From 9a3e128b948944a0131ca85fe6798186560f868b Mon Sep 17 00:00:00 2001
From: Vladimir Magamedov <vladimir@magamedov.com>
Date: Tue, 11 Jun 2019 21:29:30 +0300
Subject: [PATCH] Added simple example of the mTLS setup

---
 Makefile                    |  2 ++
 examples/mtls/__init__.py   |  0
 examples/mtls/client.py     | 50 +++++++++++++++++++++++++++++++++++++
 examples/mtls/keys/Makefile |  9 +++++++
 examples/mtls/server.py     | 50 +++++++++++++++++++++++++++++++++++++
 5 files changed, 111 insertions(+)
 create mode 100644 examples/mtls/__init__.py
 create mode 100644 examples/mtls/client.py
 create mode 100644 examples/mtls/keys/Makefile
 create mode 100644 examples/mtls/server.py

diff --git a/Makefile b/Makefile
index bfabbfc..c0e5cd4 100644
--- a/Makefile
+++ b/Makefile
@@ -1,3 +1,5 @@
+include examples/mtls/keys/Makefile
+
 __default__:
 	@echo "Please specify a target to make"
 
diff --git a/examples/mtls/__init__.py b/examples/mtls/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/examples/mtls/client.py b/examples/mtls/client.py
new file mode 100644
index 0000000..96616f7
--- /dev/null
+++ b/examples/mtls/client.py
@@ -0,0 +1,50 @@
+import os
+import ssl
+import asyncio
+import logging
+
+from pathlib import Path
+
+from grpclib.client import Channel
+from grpclib.health.v1.health_pb2 import HealthCheckRequest
+from grpclib.health.v1.health_grpc import HealthStub
+
+
+DIR = Path(__file__).parent.joinpath('keys')
+SPY_MODE = 'SPY_MODE' in os.environ
+
+SERVER_CERT = DIR.joinpath('mccoy.pem')
+CLIENT_CERT = DIR.joinpath('spock-imposter.pem' if SPY_MODE else 'spock.pem')
+CLIENT_KEY = DIR.joinpath('spock-imposter.key' if SPY_MODE else 'spock.key')
+
+
+def create_secure_context(
+    client_cert: Path, client_key: Path, *, trusted: Path,
+) -> ssl.SSLContext:
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS)
+    ctx.verify_mode = ssl.CERT_REQUIRED
+    ctx.load_cert_chain(str(client_cert), str(client_key))
+    ctx.load_verify_locations(str(trusted))
+    ctx.options |= ssl.OP_NO_TLSv1 | ssl.OP_NO_TLSv1_1
+    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
+
+
+async def main(*, host: str = 'localhost', port: int = 50051) -> None:
+    channel = Channel(host, port, ssl=create_secure_context(
+        CLIENT_CERT, CLIENT_KEY, trusted=SERVER_CERT,
+    ))
+    stub = HealthStub(channel)
+    response = await stub.Check(HealthCheckRequest())
+    print(response)
+    channel.close()
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    asyncio.run(main())
diff --git a/examples/mtls/keys/Makefile b/examples/mtls/keys/Makefile
new file mode 100644
index 0000000..9707a66
--- /dev/null
+++ b/examples/mtls/keys/Makefile
@@ -0,0 +1,9 @@
+mkfile_dir := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
+
+keys:
+	rm -f $(mkfile_dir)*.key
+	rm -f $(mkfile_dir)*.pem
+	openssl req -x509 -newkey rsa:2048 -nodes -subj '/CN=mccoy.earth.svc.cluster.local' -keyout $(mkfile_dir)mccoy.key -out $(mkfile_dir)mccoy.pem
+	openssl req -x509 -newkey rsa:2048 -nodes -subj '/CN=mccoy.earth.svc.cluster.local' -keyout $(mkfile_dir)mccoy-imposter.key -out $(mkfile_dir)mccoy-imposter.pem
+	openssl req -x509 -newkey rsa:2048 -nodes -subj '/CN=spock.vulcan.svc.cluster.local' -keyout $(mkfile_dir)spock.key -out $(mkfile_dir)spock.pem
+	openssl req -x509 -newkey rsa:2048 -nodes -subj '/CN=spock.vulcan.svc.cluster.local' -keyout $(mkfile_dir)spock-imposter.key -out $(mkfile_dir)spock-imposter.pem
diff --git a/examples/mtls/server.py b/examples/mtls/server.py
new file mode 100644
index 0000000..efd2d1a
--- /dev/null
+++ b/examples/mtls/server.py
@@ -0,0 +1,50 @@
+import os
+import ssl
+import asyncio
+import logging
+
+from pathlib import Path
+
+from grpclib.utils import graceful_exit
+from grpclib.server import Server
+from grpclib.health.service import Health
+
+
+DIR = Path(__file__).parent.joinpath('keys')
+SPY_MODE = 'SPY_MODE' in os.environ
+
+CLIENT_CERT = DIR.joinpath('spock.pem')
+SERVER_CERT = DIR.joinpath('mccoy-imposter.pem' if SPY_MODE else 'mccoy.pem')
+SERVER_KEY = DIR.joinpath('mccoy-imposter.key' if SPY_MODE else 'mccoy.key')
+
+
+def create_secure_context(
+    server_cert: Path, server_key: Path, *, trusted: Path,
+) -> ssl.SSLContext:
+    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS)
+    ctx.verify_mode = ssl.CERT_REQUIRED
+    ctx.load_cert_chain(str(server_cert), str(server_key))
+    ctx.load_verify_locations(str(trusted))
+    ctx.options |= ssl.OP_NO_TLSv1 | ssl.OP_NO_TLSv1_1
+    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
+
+
+async def main(*, host: str = '0.0.0.0', port: int = 50051) -> None:
+    server = Server([Health()])
+    with graceful_exit([server]):
+        await server.start(host, port, ssl=create_secure_context(
+            SERVER_CERT, SERVER_KEY, trusted=CLIENT_CERT,
+        ))
+        print(f'Serving on {host}:{port}')
+        await server.wait_closed()
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.INFO)
+    asyncio.run(main())