A working patch for GeoIP in CTFd

geoip-ctfd.patch

diff --git a/CTFd/auth.py b/CTFd/auth.py
index 91f7c73..410ffab 100644
--- a/CTFd/auth.py
+++ b/CTFd/auth.py
@@ -22,6 +22,8 @@ from CTFd.utils.modes import TEAMS_MODE, USERS_MODE
 from CTFd.utils.security.signing import serialize, unserialize, SignatureExpired, BadSignature, BadTimeSignature
 from CTFd.utils.helpers import info_for, error_for, get_errors, get_infos
 from CTFd.utils.config.visibility import registration_visible
+from CTFd.utils.user import get_ip
+from CTFd.utils.countries import lookup_geoip_country_code, lookup_country_code
 
 import base64
 import requests
@@ -132,6 +134,10 @@ def register():
         name = request.form['name']
         email_address = request.form['email']
         password = request.form['password']
+        country = request.form['country']
+
+        if not country:
+            country = None
 
         name_len = len(name) == 0
         names = Users.query.add_columns('name', 'id').filter_by(name=name).first()
@@ -140,6 +146,7 @@ def register():
         pass_long = len(password) > 128
         valid_email = validators.validate_email(request.form['email'])
         team_name_email_check = validators.validate_email(name)
+        country_valid = bool(lookup_country_code(country))
 
         if not valid_email:
             errors.append("Please enter a valid email address")
@@ -161,6 +168,8 @@ def register():
             errors.append('Pick a shorter password')
         if name_len:
             errors.append('Pick a longer team name')
+        if country is not None and not country_valid:
+            errors.append('You picked an invalid country')
 
         if len(errors) > 0:
             return render_template(
@@ -168,14 +177,16 @@ def register():
                 errors=errors,
                 name=request.form['name'],
                 email=request.form['email'],
-                password=request.form['password']
+                password=request.form['password'],
+                country=request.form['country']
             )
         else:
             with app.app_context():
                 user = Users(
                     name=name.strip(),
                     email=email_address.lower(),
-                    password=password.strip()
+                    password=password.strip(),
+                    country=country,
                 )
                 db.session.add(user)
                 db.session.commit()
@@ -199,7 +210,12 @@ def register():
         db.session.close()
         return redirect(url_for('challenges.listing'))
     else:
-        return render_template('register.html', errors=errors)
+        # Incentivize people to declare their country by selecting it for them!
+        geoip = lookup_geoip_country_code(get_ip())
+        if geoip and lookup_country_code(geoip.country):
+            return render_template('register.html', errors=errors, country=geoip.country)
+        else:
+            return render_template('register.html', errors=errors)
 
 
 @auth.route('/login', methods=['POST', 'GET'])
diff --git a/CTFd/schemas/users.py b/CTFd/schemas/users.py
index 9b71ef2..87e0142 100644
--- a/CTFd/schemas/users.py
+++ b/CTFd/schemas/users.py
@@ -50,8 +50,9 @@ class UserSchema(ma.ModelSchema):
     country = field_for(
         Users,
         'country',
+        required=True,
         validate=[
-            validate_country_code
+            validate_country_code,
         ]
     )
     password = field_for(
diff --git a/CTFd/themes/core/templates/register.html b/CTFd/themes/core/templates/register.html
index 937f784..8fde80e 100644
--- a/CTFd/themes/core/templates/register.html
+++ b/CTFd/themes/core/templates/register.html
@@ -29,6 +29,18 @@
 					<input class="form-control" type="text" name="name" id="name-input" {% if name %}value="{{ name }}"{% endif %} />
 				</div>
 				<div class="form-group">
+					<label for="country-input">
+						Country (GeoIP)
+					</label>
+					<select class="form-control" id="country-input" name="country">
+						<option value="">None</option>
+						{% set countries = get_countries() %}
+						{% for country_code in countries.keys() %}
+							<option value="{{ country_code }}" {% if country == country_code %}selected{% endif %}>{{ countries[country_code] }}</option>
+						{% endfor %}
+					</select>
+				</div>
+				<div class="form-group">
 					<label for="email-input">
 						Email
 					</label>
diff --git a/CTFd/themes/core/templates/settings.html b/CTFd/themes/core/templates/settings.html
index 909e061..1efd8e9 100644
--- a/CTFd/themes/core/templates/settings.html
+++ b/CTFd/themes/core/templates/settings.html
@@ -72,7 +72,7 @@
 						Country
 					</label>
 					<select class="form-control" id="country-input" name="country">
-						<option></option>
+						<option value="">None</option>
 						{% set countries = get_countries() %}
 						{% for country_code in countries.keys() %}
 							<option value="{{ country_code }}" {% if country == country_code %}selected{% endif %}>{{ countries[country_code] }}</option>
diff --git a/CTFd/utils/countries/__init__.py b/CTFd/utils/countries/__init__.py
index 5b454f3..bbe17d2 100644
--- a/CTFd/utils/countries/__init__.py
+++ b/CTFd/utils/countries/__init__.py
@@ -1,6 +1,17 @@
 # -*- coding: utf-8 -*-
 
 from collections import OrderedDict
+import os
+
+GEOIP_DB_FILE = os.path.join(os.path.dirname(__file__), 'GeoLite2-Country.mmdb')
+GEOIP_DB = None
+
+try:
+    from geoip import open_database
+    GEOIP_DB = open_database(GEOIP_DB_FILE)
+    print("Opened GeoIP database %s" % GEOIP_DB_FILE)
+except IOError:
+    print("Unable to find the GeoIP database %s" % GEOIP_DB_FILE)
 
 COUNTRIES_LIST = [
     ("AF", "Afghanistan"),
@@ -269,3 +280,9 @@ def get_countries():
 
 def lookup_country_code(country_code):
     return COUNTRIES_DICT.get(country_code)
+
+def lookup_geoip_country_code(ip_address):
+    try:
+        return GEOIP_DB.lookup(ip_address)
+    except LookupError:
+        return None