#!/usr/bin/env python
# coding: utf-8

# # End-to-end encrypted communication with a Jupyter kernel
# 
# Prompted by [this question](https://discourse.jupyter.org/t/how-do-i-properly-protect-my-data-access-passwords-not-jupyter-tokens-passwords-on-3rd-party-jupyter-hub-services/9277/2),
# 
# Disclosure: I am doing exactly what people shouldn't, which is copying and pasting from [cryptography's hazmat docs](https://cryptography.io/en/latest/hazmat/primitives/asymmetric/rsa/) without fully understanding everything involved.
# 
# We need two computers for this:
# 
# 1. the remote JupyterHub host, where we don't want credentials at rest, and don't want to trust that all network traffic is secure.
# 2. a local machine that we do trust

# Step 1: Generate a public/private key pair for sending encrypted messages to the kernel.
# 
# The private key only resides in memory, and is never displayed or shared, and it never leaves the kernel's memory.
# 
# It is only used for the lifetime of the current kernel - a new key will be regenerated on kernel restart and we need to do the whole rigamarole again.

# In[1]:


# remote
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa

private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)

pem = private_key.public_key().public_bytes(
    encoding=serialization.Encoding.PEM,
    format=serialization.PublicFormat.SubjectPublicKeyInfo
)
print(pem.decode("ascii"))


# On a local, trusted computer, encrypt your message with this public key.
# 
# You need to get the public key to the local computer and load it:
# 
# The public key itself is not sensitive, so this can be done via insecure channels:

# In[2]:


# local
from cryptography.hazmat.primitives import serialization

public_key = serialization.load_pem_public_key("""
-----BEGIN PUBLIC KEY-----
MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEA3ylCX42ZsgDYDqWG8R6c
KU6YPiHPekSy+OG8vc+crDiadwEKaZ1dakZvJWbz38r9WyCy1WA4pszChEynlqNj
g8VjeoxlWab6IwJtc44I8IqFoYmWPuKImaZ+kkRxQbZW4sIeMco7R8GBce2nK6PJ
DnfyIRJINqotXs9DVMEHgmJ/Lt39gZVibu6ZulbMUgoj/xcMRqNdlDeukjAMOrgY
smkilRtELdFZPR7Dl+8PKTOJ8KEMCapIsIkl89RxY0SuIokrrsrHROnIpSHVDow8
cyfR0/zO4CJahkO6Ymp7DzJq1+Yc8vEbWMx/7tuln+SgGFx16VxXPd5NyOWOu2Wv
7wIDAQAB
-----END PUBLIC KEY-----
""".encode('ascii'))


# Now retrieve our credentials on the local computer, however makes sense.
# 
# We need it in bytes form (could be reading a file, but I'll use a json blob)

# In[3]:


# local
import json

creds = {"key": "value"}
creds_bytes = json.dumps(creds).encode("utf8")


# Now serialize our credentials message using the kernel's public key

# In[4]:


# local
import binascii
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding

encrypted_creds = public_key.encrypt(
    creds_bytes,
    padding.OAEP(
        mgf=padding.MGF1(algorithm=hashes.SHA256()),
        algorithm=hashes.SHA256(),
        label=None,
    ),
)

encrypted_creds
b64_encrypted_creds = binascii.b2a_base64(encrypted_creds).decode('ascii')
print(b64_encrypted_creds)


# At this point, we have an encrypted form of the credentials,
# which should only be readable with the private key in memory of your remote kernel.
# 
# 
# Back to the remote, shared system.
# 
# We can input this encrypted message using `getpass()` to avoid it being at rest in your inputs or execution history.

# In[5]:


# remote
import binascii
from getpass import getpass

b64_encrypted_creds = getpass("Base64 encrypted credentials: ")
encrypted_creds = binascii.a2b_base64(b64_encrypted_creds)


# Now we can decrypt the credentials on the host:

# In[6]:


# remote
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding

creds_bytes = private_key.decrypt(
    encrypted_creds,
    padding.OAEP(
        mgf=padding.MGF1(algorithm=hashes.SHA256()),
        algorithm=hashes.SHA256(),
        label=None,
    ),
)
json.loads(creds_bytes.decode("utf8"))


# This is super tedious and I'm not sure I would recommend it to anyone, but it is an example of e2e communication with a Jupyter kernel.
# 
# Also never take security advice from me, I don't know what I'm doing.