From 69198ff08ad29e7b24e7d1e8dbcb181bc3ca13cc Mon Sep 17 00:00:00 2001
From: Son NK <son@simplelogin.io>
Date: Sat, 14 Mar 2020 22:24:02 +0100
Subject: [PATCH] delete all unnecessary headers in PGP

---
 app/email_utils.py |  9 +++++++++
 email_handler.py   | 12 ++++++++++++
 init_app.py        |  2 ++
 server.py          |  2 +-
 4 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/app/email_utils.py b/app/email_utils.py
index 2db45f8a..2099d169 100644
--- a/app/email_utils.py
+++ b/app/email_utils.py
@@ -307,6 +307,15 @@ def delete_header(msg: Message, header: str):
             del msg._headers[i]
 
 
+def delete_all_headers_except(msg: Message, headers: [str]):
+    headers = [h.lower() for h in headers]
+
+    for i in reversed(range(len(msg._headers))):
+        header_name = msg._headers[i][0].lower()
+        if header_name not in headers:
+            del msg._headers[i]
+
+
 def email_belongs_to_alias_domains(email: str) -> bool:
     """return True if an email ends with one of the alias domains provided by SimpleLogin"""
     for domain in ALIAS_DOMAINS:
diff --git a/email_handler.py b/email_handler.py
index 238510bf..c6110204 100644
--- a/email_handler.py
+++ b/email_handler.py
@@ -64,6 +64,7 @@ from app.email_utils import (
     email_belongs_to_alias_domains,
     render,
     get_orig_message_from_bounce,
+    delete_all_headers_except,
 )
 from app.extensions import db
 from app.log import LOG
@@ -269,6 +270,17 @@ def prepare_pgp_message(orig_msg: Message, pgp_fingerprint: str):
         if header_name != "Content-Type".lower():
             msg[header_name] = orig_msg._headers[i][1]
 
+    # Delete unnecessary headers in orig_msg except to save space
+    delete_all_headers_except(
+        orig_msg,
+        [
+            "MIME-Version",
+            "Content-Type",
+            "Content-Disposition",
+            "Content-Transfer-Encoding",
+        ],
+    )
+
     first = MIMEApplication(
         _subtype="pgp-encrypted", _encoder=encoders.encode_7or8bit, _data=""
     )
diff --git a/init_app.py b/init_app.py
index bf5fc728..54bb1239 100644
--- a/init_app.py
+++ b/init_app.py
@@ -22,6 +22,8 @@ def load_pgp_public_keys(app):
 
         db.session.commit()
 
+    LOG.d("Finish load_pgp_public_keys")
+
 
 if __name__ == "__main__":
     app = create_app()
diff --git a/server.py b/server.py
index 46abbff9..24c27d0f 100644
--- a/server.py
+++ b/server.py
@@ -151,7 +151,7 @@ def fake_data():
     db.session.commit()
 
     api_key = ApiKey.create(user_id=user.id, name="Chrome")
-    api_key.code = "codeCH"
+    api_key.code = "code"
 
     api_key = ApiKey.create(user_id=user.id, name="Firefox")
     api_key.code = "codeFF"