69 lines
2.3 KiB
Python
69 lines
2.3 KiB
Python
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()
|