DjangoAdmin_Flask_EleVue/utils/ip2region.py

"""
IP地理位置查询模块（基于ip2region离线数据库）
=============================================================

基于 ip2region.xdb 离线数据库提供高性能IP地理位置查询，无需网络请求。

核心功能：
- get_ip_region: 根据IP获取完整区域信息字典
- get_ip_location: 根据IP获取地理位置文本（兼容现有调用）
- get_ip_country / get_ip_province / get_ip_city / get_ip_isp: 按层级获取

xdb文件位于 utils/ip2region.xdb，启动时全量加载到内存，查询零IO延迟。
"""

import logging
import os

from ip2region import searcher, util

logger = logging.getLogger(__name__)

# xdb数据库文件路径（与本模块同目录）
_XDB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ip2region.xdb")

# 模块级单例：启动时加载xdb到内存，后续查询零IO
_searcher = None


def _get_searcher():
    """获取搜索器单例（懒加载，首次调用时初始化）"""
    global _searcher
    if _searcher is not None:
        return _searcher

    if not os.path.exists(_XDB_PATH):
        logger.error(f"ip2region数据库文件不存在: {_XDB_PATH}")
        return None

    try:
        c_buff = util.load_content_from_file(_XDB_PATH)
        _searcher = searcher.new_with_buffer(util.IPv4, c_buff)
        logger.info(f"ip2region数据库加载成功: {_XDB_PATH}")
    except Exception as e:
        logger.error(f"ip2region数据库加载失败: {e}")
        return None

    return _searcher


def _parse_region(region_str):
    """
    解析ip2region原始结果字符串

    原始格式: "国家|省份|城市|ISP|国家代码"
    示例: "中国|浙江省|杭州市|阿里|CN", "United States|California|0|Google LLC|US"

    Args:
        region_str: ip2region返回的原始字符串

    Returns:
        dict: 各层级区域信息
    """
    parts = region_str.split("|") if region_str else []
    # 补齐到5段
    while len(parts) < 5:
        parts.append("0")

    country = parts[0]
    province = parts[1]
    city = parts[2]
    isp = parts[3]
    country_code = parts[4]

    # 清理占位符 "0" → 空字符串
    province = "" if province == "0" else province
    city = "" if city == "0" else city
    isp = "" if isp == "0" else isp
    country_code = "" if country_code == "0" else country_code

    return {
        "country": country,
        "province": province,
        "city": city,
        "isp": isp,
        "country_code": country_code,
    }


def _is_private_ip(ip):
    """判断是否为内网/保留IP"""
    if not ip:
        return True
    ip = ip.strip()
    if ip in ("127.0.0.1", "localhost", "0.0.0.0"):
        return True
    if ip.startswith("10.") or ip.startswith("192.168."):
        return True
    if ip.startswith("172."):
        parts = ip.split(".")
        if len(parts) >= 2:
            second = int(parts[1])
            if 16 <= second <= 31:
                return True
    return False


def get_ip_region(ip):
    """
    根据IP获取完整区域信息字典

    Args:
        ip: IP地址字符串

    Returns:
        dict: 区域信息，包含 country/province/city/isp/country_code 字段。
              内网IP返回 {"country": "内网IP", ...}，查询失败返回 {"country": "未知", ...}
    """
    if _is_private_ip(ip):
        return {"country": "内网IP", "province": "", "city": "", "isp": "", "country_code": ""}

    s = _get_searcher()
    if s is None:
        return {"country": "未知", "province": "", "city": "", "isp": "", "country_code": ""}

    try:
        region_str = s.search(ip)
        if not region_str:
            return {"country": "未知", "province": "", "city": "", "isp": "", "country_code": ""}
        return _parse_region(region_str)
    except Exception as e:
        logger.warning(f"IP地理位置查询失败 [{ip}]: {e}")
        return {"country": "未知", "province": "", "city": "", "isp": "", "country_code": ""}


def get_ip_location(ip):
    """
    根据IP获取地理位置文本描述

    格式: "国家 省份 城市"（与原 utils/ip.py 的 get_ip_location 返回格式兼容）
    内网IP返回 "内网IP"，查询失败返回 "未知"

    Args:
        ip: IP地址字符串

    Returns:
        str: 地理位置文本
    """
    region = get_ip_region(ip)

    if region["country"] in ("内网IP", "未知"):
        return region["country"]

    parts = [region["country"]]
    if region["province"]:
        parts.append(region["province"])
    if region["city"] and region["city"] != region["province"]:
        parts.append(region["city"])

    return " ".join(parts)


def get_ip_country(ip):
    """根据IP获取国家"""
    return get_ip_region(ip)["country"]


def get_ip_province(ip):
    """根据IP获取省份"""
    return get_ip_region(ip)["province"]


def get_ip_city(ip):
    """根据IP获取城市"""
    return get_ip_region(ip)["city"]


def get_ip_isp(ip):
    """根据IP获取运营商"""
    return get_ip_region(ip)["isp"]