fix: 修复获取 city 时可能的报错 (#8295)

Co-authored-by: ibuler <ibuler@qq.com>
jumpserver · May 24, 2022 · c30b024 · c30b024
1 parent 9c14eb5
commit c30b024
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 21 deletions.
diff --git a/apps/common/utils/ip/geoip/utils.py b/apps/common/utils/ip/geoip/utils.py
@@ -13,10 +13,6 @@
 
 
 def get_ip_city_by_geoip(ip):
-    if not ip or '.' not in ip or not isinstance(ip, str):
-        return _("Invalid ip")
-    if ':' in ip:
-        return 'IPv6'
     global reader
     if reader is None:
         path = os.path.join(os.path.dirname(__file__), 'GeoLite2-City.mmdb')
@@ -32,7 +28,7 @@ def get_ip_city_by_geoip(ip):
     try:
         response = reader.city(ip)
     except GeoIP2Error:
-        return {}
+        return _("Unknown")
 
     city_names = response.city.names or {}
     lang = settings.LANGUAGE_CODE[:2]

diff --git a/apps/common/utils/ip/ipip/utils.py b/apps/common/utils/ip/ipip/utils.py
@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 #
 import os
-from django.utils.translation import ugettext as _
 
 import ipdb
 
@@ -11,13 +10,13 @@
 
 def get_ip_city_by_ipip(ip):
     global ipip_db
-    if not ip or not isinstance(ip, str):
-        return _("Invalid ip")
-    if ':' in ip:
-        return 'IPv6'
     if ipip_db is None:
         ipip_db_path = os.path.join(os.path.dirname(__file__), 'ipipfree.ipdb')
         ipip_db = ipdb.City(ipip_db_path)
-
-    info = ipip_db.find_info(ip, 'CN')
+    try:
+        info = ipip_db.find_info(ip, 'CN')
+    except ValueError:
+        return None
+    if not info:
+        raise None
     return {'city': info.city_name, 'country': info.country_name}
diff --git a/apps/common/utils/ip/utils.py b/apps/common/utils/ip/utils.py
@@ -74,13 +74,18 @@ def contains_ip(ip, ip_group):
 
 
 def get_ip_city(ip):
-    info = get_ip_city_by_ipip(ip)
-    city = info.get('city', _("Unknown"))
-    country = info.get('country')
+    if not ip or not isinstance(ip, str):
+        return _("Invalid ip")
+    if ':' in ip:
+        return 'IPv6'
 
-    # 国内城市 并且 语言是中文就使用国内
-    is_zh = settings.LANGUAGE_CODE.startswith('zh')
-    if country == '中国' and is_zh:
-        return city
-    else:
-        return get_ip_city_by_geoip(ip)
+    info = get_ip_city_by_ipip(ip)
+    if info:
+        city = info.get('city', _("Unknown"))
+        country = info.get('country')
+
+        # 国内城市 并且 语言是中文就使用国内
+        is_zh = settings.LANGUAGE_CODE.startswith('zh')
+        if country == '中国' and is_zh:
+            return city
+    return get_ip_city_by_geoip(ip)