From e44860329b32c234740f7530fd4f9605274c2e63 Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Sun, 12 Apr 2020 20:14:49 +0200
Subject: [PATCH] Make sure user cannot create more than 50 directories

---
 app/config.py                    | 3 +++
 app/dashboard/views/directory.py | 9 ++++++++-
 app/models.py                    | 5 ++++-
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/app/config.py b/app/config.py
index 23ec8511..99f052ad 100644
--- a/app/config.py
+++ b/app/config.py
@@ -57,6 +57,9 @@ except Exception:
     print("MAX_NB_EMAIL_FREE_PLAN is not set, use 5 as default value")
     MAX_NB_EMAIL_FREE_PLAN = 5
 
+# maximum number of directory a premium user can create
+MAX_NB_DIRECTORY = 50
+
 # allow to override postfix server locally
 POSTFIX_SERVER = os.environ.get("POSTFIX_SERVER", "240.0.0.1")
 
diff --git a/app/dashboard/views/directory.py b/app/dashboard/views/directory.py
index 1dba78db..7f7753a5 100644
--- a/app/dashboard/views/directory.py
+++ b/app/dashboard/views/directory.py
@@ -3,7 +3,7 @@ from flask_login import login_required, current_user
 from flask_wtf import FlaskForm
 from wtforms import StringField, validators
 
-from app.config import EMAIL_DOMAIN, ALIAS_DOMAINS
+from app.config import EMAIL_DOMAIN, ALIAS_DOMAINS, MAX_NB_DIRECTORY
 from app.dashboard.base import dashboard_bp
 from app.extensions import db
 from app.models import Directory
@@ -46,6 +46,13 @@ def directory():
                 flash("Only premium plan can add directory", "warning")
                 return redirect(url_for("dashboard.directory"))
 
+            if current_user.nb_directory() >= MAX_NB_DIRECTORY:
+                flash(
+                    f"You cannot have more than {MAX_NB_DIRECTORY} directories",
+                    "warning",
+                )
+                return redirect(url_for("dashboard.directory"))
+
             if new_dir_form.validate():
                 new_dir_name = new_dir_form.name.data.lower()
 
diff --git a/app/models.py b/app/models.py
index 1ba7d798..3201d6ea 100644
--- a/app/models.py
+++ b/app/models.py
@@ -328,7 +328,7 @@ class User(db.Model, ModelMixin, UserMixin):
         """
         sub = Subscription.get_by(user_id=self.id)
         # TODO: sub is active only if sub.next_bill_date > now
-        # due to a bug on next_bill_date, wait until next month (April 8)
+        # due to a bug on next_bill_date, wait until next month (May 8)
         # when all next_bill_date are correctly updated to add this check
 
         if sub and sub.cancelled:
@@ -353,6 +353,9 @@ class User(db.Model, ModelMixin, UserMixin):
 
         return mailboxes
 
+    def nb_directory(self):
+        return Directory.query.filter_by(user_id=self.id).count()
+
     def __repr__(self):
         return f"<User {self.id} {self.name} {self.email}>"