From 15ac7988f333b0825a3cd166b4e041c01449aeac Mon Sep 17 00:00:00 2001
From: David Duque <david@davness.dev>
Date: Sun, 6 Dec 2020 02:27:04 +0000
Subject: [PATCH] Implement key stripping

---
 management/pgp.py |  2 +-
 management/wkd.py | 37 ++++++++++++++++++++++++++++++++++---
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/management/pgp.py b/management/pgp.py
index 76865c6..5dd2927 100755
--- a/management/pgp.py
+++ b/management/pgp.py
@@ -56,7 +56,7 @@ def fork_context(f, context = default_context):
 		with tempfile.TemporaryDirectory() as tmpdir:
 			shutil.copytree(context.home_dir, f"{tmpdir}/gnupg")
 			kwargs["context"] = gpg.Context(armor=context.armor, home_dir=f"{tmpdir}/gnupg")
-			f(*args, **kwargs)
+			return f(*args, **kwargs)
 
 	return wrapped
 
diff --git a/management/wkd.py b/management/wkd.py
index 99b38c2..0ed8984 100644
--- a/management/wkd.py
+++ b/management/wkd.py
@@ -1,6 +1,6 @@
 #!/usr/local/lib/mailinabox/env/bin/python
 # WDK (Web Key Directory) Manager: Facilitates discovery of keys by third-parties
-# Current relevant documents: https://tools.ietf.org/id/draft-koch-openpgp-webkey-service-10.html
+# Current relevant documents: https://tools.ietf.org/id/draft-koch-openpgp-webkey-service-11.html
 
 import pgp, utils
 from cryptography.hazmat.primitives import hashes
@@ -29,8 +29,39 @@ def zbase32(digest):
 
 @pgp.fork_context
 # Strips and exports a key so that only the provided UID index(es) remain.
-def key_stripped(key, uid_index, context = None):
-	pass
+# This is to comply with the following requirement, set forth in section 5 of the draft:
+#
+# The mail provider MUST make sure to publish a key in a way
+# that only the mail address belonging to the requested user
+# is part of the User ID packets included in the returned key.
+# Other User ID packets and their associated binding signatures
+# MUST be removed before publication.
+def strip_and_export(fpr, except_uid_indexes, context):
+	context.armor = False # We need to disable armor output for this key
+	k = pgp.get_key(fpr, context)
+	if k is None:
+		return None
+	for i in except_uid_indexes:
+		if i > len(k.uids):
+			raise ValueError(f"UID index {i} out of bounds")
+
+	switch = [(f"uid {i+1}" if i + 1 not in except_uid_indexes else "") for i in range(len(k.uids))] + ["deluid", "save"]
+	stage = [-1] # Horrible hack: Because it's a reference (aka pointer), we can pass these around
+	def interaction(request, prompt):
+		print(f"{request}/{prompt}")
+		if request in ["GOT_IT", "KEY_CONSIDERED", ""]:
+			return 0
+		elif request == "GET_BOOL":
+			# No way to confirm interactively, so we just say yes
+			return "y" # Yeah, I'd also rather just return True but that doesn't work
+		elif request == "GET_LINE" and prompt == "keyedit.prompt":
+			stage[0] += 1
+			return switch[stage[0]]
+		else:
+			raise Exception("No idea of what to do!")
+
+	context.interact(k, interaction)
+	return pgp.export_key(fpr, context)
 
 def set_wkd_published(fingerprint, publish):
     if pgp.get_key(fingerprint) is None: