From 0a04f0f35146fdd1bde364ddf9e6cd66f4d768e4 Mon Sep 17 00:00:00 2001
From: ibuler <ibuler@qq.com>
Date: Fri, 20 May 2022 00:10:22 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E4=B8=8B=E8=BD=BD=20ip=20=E6=95=B0?=
 =?UTF-8?q?=E6=8D=AE=E5=BA=93?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/common/utils/file.py | 10 ++++++++++
 jms                       | 17 +++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/apps/common/utils/file.py b/apps/common/utils/file.py
index c00293de5..9529fbdca 100644
--- a/apps/common/utils/file.py
+++ b/apps/common/utils/file.py
@@ -1,6 +1,8 @@
 import os
 import csv
+
 import pyzipper
+import requests
 
 
 def create_csv_file(filename, headers, rows, ):
@@ -18,3 +20,11 @@ def encrypt_and_compress_zip_file(filename, secret_password, encrypted_filenames
         for encrypted_filename in encrypted_filenames:
             with open(encrypted_filename, 'rb') as f:
                 zf.writestr(os.path.basename(encrypted_filename), f.read())
+
+
+def download_file(src, path):
+    with requests.get(src, stream=True) as r:
+        r.raise_for_status()
+        with open(path, 'wb') as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
diff --git a/jms b/jms
index 517fb0602..5d5b5c433 100755
--- a/jms
+++ b/jms
@@ -8,6 +8,7 @@ import time
 import argparse
 import sys
 import django
+import requests
 from django.core import management
 from django.db.utils import OperationalError
 
@@ -33,6 +34,7 @@ except ImportError as e:
 
 try:
     from jumpserver.const import CONFIG
+    from common.utils.file import download_file
 except ImportError as e:
     print("Import error: {}".format(e))
     print("Could not find config file, `cp config_example.yml config.yml`")
@@ -105,6 +107,20 @@ def compile_i18n_file():
     logging.info("Compile i18n files done")
 
 
+def download_ip_db():
+    db_base_dir = os.path.join(APP_DIR, 'common', 'utils', 'ip')
+    db_path_url_mapper = {
+        ('geoip', 'GeoLite2-City.mmdb'): 'https://jms-pkg.oss-cn-beijing.aliyuncs.com/ip/GeoLite2-City.mmdb',
+        ('ipip', 'ipipfree.ipdb'): 'https://jms-pkg.oss-cn-beijing.aliyuncs.com/ip/ipipfree.ipdb'
+    }
+    for p, src in db_path_url_mapper.items():
+        path = os.path.join(db_base_dir, *p)
+        if os.path.isfile(path) and os.path.getsize(path) > 1000:
+            continue
+        print("Download ip db: {}".format(path))
+        download_file(src, path)
+
+
 def upgrade_db():
     collect_static()
     perform_db_migrate()
@@ -114,6 +130,7 @@ def prepare():
     check_database_connection()
     upgrade_db()
     expire_caches()
+    download_ip_db()
 
 
 def start_services():