From cf699aede04fe483352cf169b70f87c2e5518b06 Mon Sep 17 00:00:00 2001
From: Kevin Chung <kchung@ctfd.io>
Date: Sun, 28 Jun 2020 14:09:48 -0400
Subject: [PATCH] Lookup the respective country for IP addresses in the admin
 panel (#1512)

* Adds support for looking up countries by IP address.
* Shows country and flag for a user/team's IP addresses in the admin panel
* Adds support for `GEOIP_DATABASE_PATH` in config.py
* Closes #1115
---
 .../admin/templates/modals/teams/addresses.html   | 12 ++++++++++++
 .../admin/templates/modals/users/addresses.html   | 12 ++++++++++++
 CTFd/themes/admin/templates/teams/team.html       |  2 +-
 CTFd/themes/admin/templates/users/user.html       |  2 +-
 CTFd/utils/countries/geoip.py                     | 15 +++++++++++++++
 CTFd/utils/initialization/__init__.py             |  4 +++-
 populate.py                                       | 10 ++++++++++
 requirements.txt                                  |  2 ++
 8 files changed, 56 insertions(+), 3 deletions(-)
 create mode 100644 CTFd/utils/countries/geoip.py

diff --git a/CTFd/themes/admin/templates/modals/teams/addresses.html b/CTFd/themes/admin/templates/modals/teams/addresses.html
index b3e4760d..aacd8a2d 100644
--- a/CTFd/themes/admin/templates/modals/teams/addresses.html
+++ b/CTFd/themes/admin/templates/modals/teams/addresses.html
@@ -6,6 +6,7 @@
                     <td class="text-center"><b>User</b></td>
                     <td class="text-center"><b>IP Address</b></td>
                     <td class="text-center"><b>Last Seen</b></td>
+                    <td class="text-center"><b>Country</b></td>
                 </tr>
             </thead>
             <tbody>
@@ -20,6 +21,17 @@
                     <td class="text-center solve-time">
                         <span data-time="{{ addr.date | isoformat }}"></span>
                     </td>
+                    {% set country = lookup_ip_address(addr.ip) %}
+                    {% if country %}
+                    <td class="text-center">
+                        <i class="flag-{{ country.lower() }}"></i>
+                        &nbsp;
+                        {{ lookup_country_code(country) }}
+                    </td>
+                    {% else %}
+                    <td class="text-center">
+                    </td>
+                    {% endif %}
                 </tr>
             {% endfor %}
             </tbody>
diff --git a/CTFd/themes/admin/templates/modals/users/addresses.html b/CTFd/themes/admin/templates/modals/users/addresses.html
index 218a64ee..75af0a26 100644
--- a/CTFd/themes/admin/templates/modals/users/addresses.html
+++ b/CTFd/themes/admin/templates/modals/users/addresses.html
@@ -5,6 +5,7 @@
 				<tr>
 					<td class="text-center"><b>IP Address</b></td>
 					<td class="text-center"><b>Last Seen</b></td>
+					<td class="text-center"><b>Country</b></td>
 				</tr>
 			</thead>
 			<tbody>
@@ -14,6 +15,17 @@
 					<td class="text-center solve-time">
 						<span data-time="{{ addr.date | isoformat }}"></span>
 					</td>
+					{% set country = lookup_ip_address(addr.ip) %}
+					{% if country %}
+					<td class="text-center">
+						<i class="flag-{{ country.lower() }}"></i>
+						&nbsp;
+						{{ lookup_country_code(country) }}
+					</td>
+					{% else %}
+					<td class="text-center">
+					</td>
+					{% endif %}
 				</tr>
 			{% endfor %}
 			</tbody>
diff --git a/CTFd/themes/admin/templates/teams/team.html b/CTFd/themes/admin/templates/teams/team.html
index 62cef227..dcb627ec 100644
--- a/CTFd/themes/admin/templates/teams/team.html
+++ b/CTFd/themes/admin/templates/teams/team.html
@@ -71,7 +71,7 @@
 </template>
 
 <div id="team-addresses-modal" class="modal fade">
-	<div class="modal-dialog">
+	<div class="modal-dialog modal-lg">
 		<div class="modal-content">
 			<div class="modal-header">
 				<h2 class="modal-action text-center w-100">IP Addresses</h2>
diff --git a/CTFd/themes/admin/templates/users/user.html b/CTFd/themes/admin/templates/users/user.html
index 4b48e5c9..7ab5f339 100644
--- a/CTFd/themes/admin/templates/users/user.html
+++ b/CTFd/themes/admin/templates/users/user.html
@@ -53,7 +53,7 @@
 </div>
 
 <div id="user-addresses-modal" class="modal fade">
-	<div class="modal-dialog">
+	<div class="modal-dialog modal-lg">
 		<div class="modal-content">
 			<div class="modal-header">
 				<h2 class="modal-action text-center w-100">IP Addresses</h2>
diff --git a/CTFd/utils/countries/geoip.py b/CTFd/utils/countries/geoip.py
new file mode 100644
index 00000000..f2e29c20
--- /dev/null
+++ b/CTFd/utils/countries/geoip.py
@@ -0,0 +1,15 @@
+import geoacumen
+import maxminddb
+from flask import current_app
+
+IP_ADDR_LOOKUP = maxminddb.open_database(
+    current_app.config.get("GEOIP_DATABASE_PATH", geoacumen.db_path)
+)
+
+
+def lookup_ip_address(addr):
+    response = IP_ADDR_LOOKUP.get(addr)
+    try:
+        return response["country"]["iso_code"]
+    except KeyError:
+        return None
diff --git a/CTFd/utils/initialization/__init__.py b/CTFd/utils/initialization/__init__.py
index 99f4ed88..1b61d168 100644
--- a/CTFd/utils/initialization/__init__.py
+++ b/CTFd/utils/initialization/__init__.py
@@ -20,7 +20,6 @@ from CTFd.utils.config import (
     is_setup,
 )
 from CTFd.utils.config.pages import get_pages
-from CTFd.utils.countries import get_countries, lookup_country_code
 from CTFd.utils.dates import isoformat, unix_time, unix_time_millis
 from CTFd.utils.events import EventManager, RedisEventManager
 from CTFd.utils.humanize.words import pluralize
@@ -65,6 +64,8 @@ def init_template_globals(app):
         registration_visible,
         scores_visible,
     )
+    from CTFd.utils.countries import get_countries, lookup_country_code
+    from CTFd.utils.countries.geoip import lookup_ip_address
 
     app.jinja_env.globals.update(config=config)
     app.jinja_env.globals.update(get_pages=get_pages)
@@ -85,6 +86,7 @@ def init_template_globals(app):
     app.jinja_env.globals.update(generate_account_url=generate_account_url)
     app.jinja_env.globals.update(get_countries=get_countries)
     app.jinja_env.globals.update(lookup_country_code=lookup_country_code)
+    app.jinja_env.globals.update(lookup_ip_address=lookup_ip_address)
     app.jinja_env.globals.update(accounts_visible=accounts_visible)
     app.jinja_env.globals.update(challenges_visible=challenges_visible)
     app.jinja_env.globals.update(registration_visible=registration_visible)
diff --git a/populate.py b/populate.py
index 07cd6398..af669a65 100644
--- a/populate.py
+++ b/populate.py
@@ -17,6 +17,7 @@ from CTFd.models import (
     ChallengeFiles,
     Fails,
     Solves,
+    Tracking,
 )
 from faker import Faker
 
@@ -108,6 +109,10 @@ def gen_file():
     return fake.file_name()
 
 
+def gen_ip():
+    return fake.ipv4()
+
+
 def random_date(start, end):
     return start + datetime.timedelta(
         seconds=random.randint(0, int((end - start).total_seconds()))
@@ -198,6 +203,11 @@ if __name__ == "__main__":
                     if mode == "teams":
                         user.team_id = random.randint(1, TEAM_AMOUNT)
                     db.session.add(user)
+                    db.session.flush()
+
+                    track = Tracking(ip=gen_ip(), user_id=user.id)
+                    db.session.add(track)
+                    db.session.flush()
                     count += 1
                 except Exception:
                     pass
diff --git a/requirements.txt b/requirements.txt
index 1aeb3e98..7b151f94 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -26,3 +26,5 @@ pydantic==1.5.1
 lxml==4.5.1
 html5lib==1.0.1
 WTForms==2.3.1
+python-geoacumen==0.0.1
+maxminddb==1.5.4