Closes #9

CTFd/admin.py

@@ -54,10 +54,10 @@ def init_admin(app):
 
             try:
                 view_challenges_unregistered = bool(request.form.get('view_challenges_unregistered', None))
+                prevent_registration = bool(request.form.get('prevent_registration', None))
             except (ValueError, TypeError):
                 view_challenges_unregistered = None
-
-            print repr(start), repr(end), repr(view_challenges_unregistered)
+                prevent_registration = None
 
             db_start = Config.query.filter_by(key='start').first()
             db_start.value = start
@@ -68,9 +68,13 @@ def init_admin(app):
             db_view_challenges_unregistered = Config.query.filter_by(key='view_challenges_unregistered').first()
             db_view_challenges_unregistered.value = view_challenges_unregistered
 
+            db_prevent_registration = Config.query.filter_by(key='prevent_registration').first()
+            db_prevent_registration.value = prevent_registration
+
             db.session.add(db_start)
             db.session.add(db_end)
             db.session.add(db_view_challenges_unregistered)
+            db.session.add(db_prevent_registration)
 
             db.session.commit()
             return redirect('/admin/config')
@@ -96,10 +100,18 @@ def init_admin(app):
             view_challenges_unregistered = Config('view_challenges_unregistered', None)
             db.session.add(view_challenges_unregistered)
 
+        prevent_registration = Config.query.filter_by(key='prevent_registration').first()
+        if prevent_registration:
+            prevent_registration = (prevent_registration.value == '1')
+        else:
+            prevent_registration = Config('prevent_registration', None)
+            db.session.add(prevent_registration)
+
         db.session.commit()
         db.session.close()
 
-        return render_template('admin/config.html', start=start, end=end, view_challenges_unregistered=view_challenges_unregistered)
+        return render_template('admin/config.html', start=start, end=end, view_challenges_unregistered=view_challenges_unregistered, 
+            prevent_registration=prevent_registration)
 
     @app.route('/admin/pages', defaults={'route': None}, methods=['GET', 'POST'])
     @app.route('/admin/pages/<route>', methods=['GET', 'POST'])

CTFd/auth.py

@@ -1,5 +1,5 @@
 from flask import render_template, request, redirect, abort, jsonify, url_for, session
-from CTFd.utils import sha512, is_safe_url, authed, mailserver, sendmail
+from CTFd.utils import sha512, is_safe_url, authed, mailserver, sendmail, can_register
 from CTFd.models import db, Teams
 
 from itsdangerous import TimedSerializer, BadTimeSignature
@@ -56,6 +56,8 @@ Did you initiate a password reset?
 
     @app.route('/register', methods=['POST', 'GET'])
     def register():
+        if not can_register():
+            return redirect('/login')
         if request.method == 'POST':
             errors = []
             name_len = len(request.form['name']) == 0

CTFd/utils.py

@@ -18,18 +18,14 @@ def init_utils(app):
     app.jinja_env.filters['unix_time_millis'] = unix_time_millis
     app.jinja_env.filters['long2ip'] = long2ip
     app.jinja_env.globals.update(pages=pages)
+    app.jinja_env.globals.update(can_register=can_register)
 
 def pages():
     pages = Pages.query.filter(Pages.route!="index").all()
     return pages
 
 def authed():
-    try:
-        if session['id']:
-            return True
-    except KeyError:
-        pass
-    return False
+    return bool(session.get('id', False))
 
 def is_setup():
     setup = Config.query.filter_by(key='setup').first()
@@ -44,6 +40,13 @@ def is_admin():
     else:
         return False
 
+def can_register():
+    config = Config.query.filter_by(key='prevent_registration').first()
+    if config:
+        return config.value != '1'
+    else:
+        return True
+
 def admins_only(f):
     @wraps(f)
     def decorated_function(*args, **kwargs):
@@ -84,7 +87,11 @@ def ctftime():
     return False
 
 def can_view_challenges():
-    return authed() or (Config.query.filter_by(key="view_challenges_unregistered").first().value == '1');
+    config = Config.query.filter_by(key="view_challenges_unregistered").first()
+    if config:
+        return authed() or config.value == '1'
+    else:
+        return authed()
 
 def unix_time(dt):
     epoch = datetime.datetime.utcfromtimestamp(0)

CTFd/views.py

@@ -66,6 +66,9 @@ def init_views(app):
                 ## Challenges cannot be viewed by unregistered users
                 view_challenges_unregistered = Config('view_challenges_unregistered', None)
 
+                ## Allow/Disallow registration
+                prevent_registration = Config('prevent_registration', None)
+
                 setup = Config('setup', True)
 
                 db.session.add(admin)
@@ -73,6 +76,7 @@ def init_views(app):
                 db.session.add(start)
                 db.session.add(end)
                 db.session.add(view_challenges_unregistered)
+                db.session.add(prevent_registration)
                 db.session.add(setup)
                 db.session.commit()
                 app.setup = False

templates/admin/config.html

@@ -22,7 +22,12 @@
             <label for="view_challenges_unregistered">Unregistered users can view challenges</label>
         </div>
 
-        <button class="radius" type='submit'>Submit</button>
+        <div class="row">
+            <input id="prevent_registration" name="prevent_registration" type="checkbox" {% if prevent_registration %}checked{% endif %}>
+            <label for="prevent_registration">Prevent public registration</label>
+        </div>
+
+        <button class="radius" type='submit'>Update</button>
     </form>
 </div>
 {% endblock %}

templates/base.html

@@ -37,6 +37,7 @@
             </li>
             {%else %}
             <li class="has-form">
+              {% if can_register() %} 
               <li class="has-dropdown">
                 <a href="/register">Register</a>
                 <ul class="dropdown">
@@ -48,7 +49,7 @@
                   </form>
                 </ul>
               </li>
-
+              {% endif %}
               <li class="has-dropdown">
                 <a href="/login">Login</a>
                 <ul class="dropdown">