from mpl_toolkits.basemap import Basemap
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
import matplotlib as mpl
from collections import defaultdict
import geoip2.database

# Function to read IPs from file
def read_ips_from_file(file_path):
    with open(file_path, 'r') as f:
        ips = f.read().strip().splitlines()
    return ips

# Function to geolocate IPs and count IPs per country
def geolocate_and_count(ips):
    reader = geoip2.database.Reader('GeoLite2-Country.mmdb')
    country_count = defaultdict(int)

    for ip in ips:
        try:
            response = reader.country(ip)
            country_name = response.country.name
            country_count[country_name] += 1
        except geoip2.errors.AddressNotFoundError:
            continue

    reader.close()
    return country_count

# Function to generate heatmap based on country IP counts
def generate_heatmap(country_counts, title):
    plt.figure(figsize=(16, 10))  # Increase figure size as needed
    m = Basemap(projection='mill', llcrnrlat=-60, urcrnrlat=90, llcrnrlon=-180, urcrnrlon=180, resolution='c')
    m.drawcoastlines()
    m.drawcountries()

    # Normalize counts for colormap
    max_count = max(country_counts.values())
    norm = mpl.colors.Normalize(vmin=0, vmax=max_count)
    cmap = plt.cm.get_cmap('YlOrRd')

    for country, count in country_counts.items():
        try:
            country_info = m.readshapefile('shapefiles/TM_WORLD_BORDERS-0.3', 'world', drawbounds=True)
            country_code = country_info['world_info'][country_info['world_info']['NAME'] == country]['ISO3'][0]
            if country_code != '-99':
                color = cmap(norm(count))
                poly = Polygon(m.world[country_code], facecolor=color, edgecolor='black')
                plt.gca().add_patch(poly)
        except:
            continue

    plt.title(title)
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    cbar = plt.colorbar(sm, orientation='vertical', shrink=0.7)  # Adjust shrink parameter if necessary
    cbar.set_label('Number of IPs')
    plt.show()


# Main function to execute the program
def main():
    ips = read_ips_from_file('sample.txt')
    country_counts = geolocate_and_count(ips)
    generate_heatmap(country_counts, 'Heatmap of IPs by Country')

if __name__ == "__main__":
    main()