import re
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
from geoip2.database import Reader


# Function to extract IP addresses from text using regex
def extract_ip_addresses(text):
    ip_pattern = r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}'
    return re.findall(ip_pattern, text)


# Function to check if an IP address is private
def is_private_ip(ip):
    octets = ip.split('.')
    first_octet = int(octets[0])
    second_octet = int(octets[1])

    # Check for private IP ranges
    if (first_octet == 10) or (first_octet == 172 and 16 <= second_octet <= 31) or (
            first_octet == 192 and second_octet == 168):
        return True
    else:
        return False


# Function to geolocate multiple IP addresses using GeoIP2
def geolocate_ips(ip_list):
    latitudes = []
    longitudes = []
    if ip_list:
        try:
            reader = Reader('GeoLite2-City.mmdb')  # Replace with your actual GeoIP2 database path
            for ip in ip_list:
                response = reader.city(ip)
                latitudes.append(response.location.latitude)
                longitudes.append(response.location.longitude)
            return latitudes, longitudes
        except Exception as e:
            print(f"Error geolocating IPs: {e}")
            return [], []
        finally:
            if 'reader' in locals():
                reader.close()
    else:
        return [], []


# Function to generate heatmap
def generate_heatmap(latitude_list, longitude_list, title):
    plt.figure(figsize=(50, 40))
    m = Basemap(projection='merc', llcrnrlat=-60, urcrnrlat=90, llcrnrlon=-180, urcrnrlon=180, resolution='c')

    if latitude_list and longitude_list and len(latitude_list) == len(longitude_list):
        valid_indices = [i for i in range(len(latitude_list)) if latitude_list[i] is not None and longitude_list[i] is not None]
        if valid_indices:
            latitudes = [latitude_list[i] for i in valid_indices]
            longitudes = [longitude_list[i] for i in valid_indices]

            x, y = m(longitudes, latitudes)
            m.scatter(x, y, marker='o', color='r', s=50, zorder=10)

    m.drawcoastlines()
    m.drawcountries()
    m.drawmapboundary()

    plt.title(title)
    plt.show()



# Read the contents of the text file
file_path = 'sample.txt'  # Replace with your actual file path
with open(file_path, 'r') as file:
    data = file.read()

# Extract all source and destination IP addresses
ips = extract_ip_addresses(data)
source_ips = [ip for ip in ips[::2] if not is_private_ip(ip)]  # Exclude private IPs from source list
destination_ips = [ip for ip in ips[1::2] if not is_private_ip(ip)]  # Exclude private IPs from destination list

print(f"Source IPs: {source_ips}")
print(f"Destination IPs: {destination_ips}")

# Geolocate all source and destination IPs
source_lats, source_lons = geolocate_ips(source_ips)
destination_lats, destination_lons = geolocate_ips(destination_ips)

# Generate heatmaps for source and destination IPs
generate_heatmap(source_lats, source_lons, 'Source IPs Geolocation Heatmap (excluding private)')
generate_heatmap(destination_lats, destination_lons, 'Destination IPs Geolocation Heatmap (excluding private)')