From 66a285e9fd21da0ed3ab508c77dfa4d5fffe050d Mon Sep 17 00:00:00 2001
From: Son NK <son@simplelogin.io>
Date: Thu, 19 Mar 2020 19:15:42 +0100
Subject: [PATCH] Only premium user can add PGP key

---
 .../templates/dashboard/mailbox_detail.html   | 23 +++++++++++++++----
 app/dashboard/views/mailbox_detail.py         |  7 ++++++
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/app/dashboard/templates/dashboard/mailbox_detail.html b/app/dashboard/templates/dashboard/mailbox_detail.html
index 6c306136..7af28de2 100644
--- a/app/dashboard/templates/dashboard/mailbox_detail.html
+++ b/app/dashboard/templates/dashboard/mailbox_detail.html
@@ -72,18 +72,33 @@
           <div class="card-title">
             Pretty Good Privacy (PGP)
             <div class="small-text">
-              By importing your PGP Public Key into SimpleLogin, all emails sent to {{mailbox.email}} are <b>encrypted</b> with your key.
+              By importing your PGP Public Key into SimpleLogin, all emails sent to {{ mailbox.email }} are
+              <b>encrypted</b> with your key.
             </div>
           </div>
 
+          {% if not current_user.is_premium() %}
+            <div class="alert alert-danger" role="alert">
+              This feature is only available in premium plan.
+            </div>
+          {% endif %}
+
           <div class="form-group">
             <label class="form-label">PGP Public Key</label>
 
-            <textarea name="pgp" class="form-control" rows=10 placeholder="-----BEGIN PGP PUBLIC KEY BLOCK-----">{{mailbox.pgp_public_key or ""}}</textarea>
+            <textarea name="pgp"
+                {% if not current_user.is_premium() %} disabled {% endif %}
+                      class="form-control" rows=10
+                      placeholder="-----BEGIN PGP PUBLIC KEY BLOCK-----">{{ mailbox.pgp_public_key or "" }}</textarea>
           </div>
 
-          <button class="btn btn-primary" name="action" value="save">Save</button>
-          <button class="btn btn-danger float-right" name="action" value="remove">Remove</button>
+          <button class="btn btn-primary" name="action"
+              {% if not current_user.is_premium() %} disabled {% endif %}
+                  value="save">Save
+          </button>
+          {% if mailbox.pgp_finger_print %}
+            <button class="btn btn-danger float-right" name="action" value="remove">Remove</button>
+          {% endif %}
 
         </div>
       </form>
diff --git a/app/dashboard/views/mailbox_detail.py b/app/dashboard/views/mailbox_detail.py
index 93975342..c74baf06 100644
--- a/app/dashboard/views/mailbox_detail.py
+++ b/app/dashboard/views/mailbox_detail.py
@@ -101,6 +101,12 @@ def mailbox_detail_route(mailbox_id):
                     )
         elif request.form.get("form-name") == "pgp":
             if request.form.get("action") == "save":
+                if not current_user.is_premium():
+                    flash("Only premium plan can add PGP Key", "warning")
+                    return redirect(
+                        url_for("dashboard.mailbox_detail_route", mailbox_id=mailbox_id)
+                    )
+
                 mailbox.pgp_public_key = request.form.get("pgp")
                 try:
                     mailbox.pgp_finger_print = load_public_key(mailbox.pgp_public_key)
@@ -113,6 +119,7 @@ def mailbox_detail_route(mailbox_id):
                         url_for("dashboard.mailbox_detail_route", mailbox_id=mailbox_id)
                     )
             elif request.form.get("action") == "remove":
+                # Free user can decide to remove their added PGP key
                 mailbox.pgp_public_key = None
                 mailbox.pgp_finger_print = None
                 db.session.commit()