import geoip2.database
import matplotlib
import matplotlib.pyplot as plt
import geopandas as gpd
import pandas as pd
import os
import sys
from collections import defaultdict
import matplotlib.font_manager as fm
import warnings

# 配置参数
IP_FILE = 'malicious_ips.txt'
GEOIP_DB = 'GeoLite2-City.mmdb'
OUTPUT_CHART = 'ip_geo_distribution.png'
MAP_DATA_URL = 'https://naciscdn.org/naturalearth/50m/cultural/ne_50m_admin_0_countries.zip'

# 内存泄漏修复
os.environ["OMP_NUM_THREADS"] = "1"

# 中文字体配置
FONT_PATH = 'C:/Windows/Fonts/msyh.ttc' if os.name == 'nt' else \
            '/System/Library/Fonts/Supplemental/Songti.ttc'
if not os.path.exists(FONT_PATH):
    FONT_PATH = fm.findfont(fm.FontProperties(family=['sans-serif']))

def configure_matplotlib():
    """跨平台字体配置"""
    try:
        plt.rcParams['font.sans-serif'] = fm.FontProperties(fname=FONT_PATH).get_name()
        plt.rcParams['axes.unicode_minus'] = False
    except Exception as e:
        warnings.warn(f"字体配置异常: {str(e)}")

def version_aware_legend():
    """版本兼容的图例参数生成"""
    try:
        from packaging.version import Version
        mpl_ver = Version(matplotlib.__version__)
        return {'title': '恶意IP数量'} if mpl_ver >= Version('3.4') else {'label': '恶意IP数量'}
    except ImportError:
        print("建议安装packaging库以获得更好的版本兼容性：pip install packaging")
        return {'label': '恶意IP数量'}

def visualize_distribution(data):
    """最终稳定版地理可视化"""
    try:
        configure_matplotlib()
        world = gpd.read_file(MAP_DATA_URL)
        
        # 国家名称标准化
        name_mapping = {
            'United States of America': 'United States',
            'Russian Federation': 'Russia',
            'Iran (Islamic Republic of)': 'Iran',
            'Viet Nam': 'Vietnam',
            'Korea, Republic of': 'South Korea',
            'Hong Kong S.A.R.': 'Hong Kong',
            'Taiwan': 'Taiwan Province of China'
        }
        world['NAME'] = world['NAME'].replace(name_mapping)
        
        # 合并数据
        df = world.merge(
            pd.DataFrame.from_dict(data, orient='index', columns=['count']),
            how="left",
            left_on='NAME',
            right_index=True
        )
        
        # 创建绘图画布
        fig, ax = plt.subplots(figsize=(20, 15))
        
        # 绘图参数配置
        plot_params = {
            'column': 'count',
            'ax': ax,
            'cmap': 'YlOrRd',
            'edgecolor': 'black',
            'linewidth': 0.3,
            'missing_kwds': {"color": "lightgrey"},
            'legend': True,
            'legend_kwds': version_aware_legend()
        }
        
        # 智能分类方案
        try:
            import mapclassify
            plot_params['scheme'] = 'NaturalBreaks'
        except ImportError:
            plot_params['scheme'] = 'equal_interval'
        
        # 绘制地图
        df.plot(**plot_params)
        
        # 添加国家标签
        top_countries = sorted(data.items(), key=lambda x: x[1], reverse=True)[:15]
        for country, count in top_countries:
            try:
                geom = df[df['NAME'] == country].geometry
                if not geom.empty:
                    centroid = geom.centroid
                    ax.text(
                        x=centroid.x.values[0],
                        y=centroid.y.values[0],
                        s=f"{country}\n{count}",
                        fontproperties=fm.FontProperties(fname=FONT_PATH, size=8),
                        ha='center',
                        va='center',
                        bbox=dict(facecolor='white', alpha=0.7, edgecolor='none')
                    )
            except Exception as e:
                continue
        
        plt.title('全球恶意IP分布热力图', 
                 fontproperties=fm.FontProperties(fname=FONT_PATH, size=18))
        plt.axis('off')
        plt.savefig(OUTPUT_CHART, dpi=400, bbox_inches='tight')
        plt.close()
        
    except Exception as e:
        print(f"地图渲染失败: {str(e)}")
        if "legend_kwds" in str(e):
            print("解决方案：")
            print("1. 升级matplotlib: pip install matplotlib --upgrade")
            print("2. 或修改代码中version_aware_legend()的返回值")

if __name__ == '__main__':
    warnings.filterwarnings("ignore", category=UserWarning)
    
    # 主流程
    ips = [ip.strip() for ip in open(IP_FILE) if ip.strip()]
    print(f"成功加载 {len(ips)} 个IP地址")
    
    # 实际地理查询
    geo_data = defaultdict(int)
    with geoip2.database.Reader(GEOIP_DB) as reader:
        for ip in ips:
            try:
                response = reader.city(ip)
                if response.country.name:
                    geo_data[response.country.name] += 1
            except:
                continue
    
    print("地理位置统计：")
    for country, count in sorted(geo_data.items(), key=lambda x: x[1], reverse=True)[:15]:
        print(f"{country}: {count}")
    
    visualize_distribution(geo_data)
    print(f"最终图表已保存至 {OUTPUT_CHART}")