#!/usr/bin/perl

use strict;
use warnings;
use Tk;
use Imager; # <-- ADDED: For robust image handling
use POSIX qw(tan);
use constant PI => 3.14159265;
use Fcntl qw(SEEK_SET);
use Net::IP;
use File::Basename;
use Text::CSV;

# --- Configuration ---
my $home_dir = $ENV{'HOME'} or die "HOME environment variable not set";
my $resource_dir = "resources";
my $image_path = "$resource_dir/w1000-old.png";
my $ipv4_db_path = "$resource_dir/ipv4.csv";
my $lsof_sudo_required_on = 1; # to get process names
my $geolite_file = "$resource_dir/GeoLite2-ASN-Blocks-IPv4.csv";

# --- Global State ---
my $state = {
    current_points => [],
    map_width      => 1000,
    map_height     => 500,
    dot_diameter   => 8,
    # Colors are now stored in Tk's hexadecimal format '#rrggbb'
    default_dot_color => '#ff3333', # red
    http_dot_color    => '#3380ff', # blue
    irc_dot_color     => '#ffff33', # yellow
    ssh_dot_color     => '#9933cc', # purple
    imap_dot_color    => '#1ae633', # green
    shadsocks_dot_color => '#808080', # gray
    text_color        => '#ffffff', # white
    do_asnlookup      => 0,
    ip2asn            => {},
};

#======================================================================
# IP Geolocation and Coordinate Functions (RAM-Efficient)
# (These functions are independent of the GUI toolkit and remain unchanged)
#======================================================================

sub ip_to_decimal {
    my ($ip_str) = @_;
    my @parts = split /\./, $ip_str;
    return 0 unless @parts == 4;
    return ($parts[0] << 24) + ($parts[1] << 16) + ($parts[2] << 8) + $parts[3];
}

sub find_location_on_disk {
    my ($fh, $ip_decimal, $file_size) = @_;
    my $low_pos = 0;
    my $high_pos = $file_size;
    for (my $i = 0; $i < 30; $i++) {
        last if ($high_pos - $low_pos) < 2;
        my $mid_pos = int($low_pos + ($high_pos - $low_pos) / 2);
        seek($fh, $mid_pos, 0);
        <$fh>;
        my $line = <$fh>;
        return undef unless defined $line;
        chomp $line;
        my ($start, $end, $lat, $lon) = split ',', $line;
        next unless defined $start and defined $end;
        if ($ip_decimal >= $start && $ip_decimal <= $end) {
            return { lat => $lat, lon => $lon };
        }
        if ($ip_decimal < $start) {
            $high_pos = $mid_pos;
        } else {
            $low_pos = $mid_pos;
        }
    }
    return undef;
}

sub get_connections {
    # Execute the ss command to get the full TCP connection details for established IPv4 sockets.
    my $ss_cmd = q(ss -t -n -4 state established);
    my $output = `$ss_cmd`;

    # Split the output into individual lines.
    my @lines = split /\n/, $output;

    # Remove the header line from the output.
    shift @lines;
    #print scalar @lines;

    my @connections;
    for my $line (@lines) {
        # Split the line into fields based on whitespace.
        # The local address:port is the 4th field and the peer address:port is the 5th.
        my @fields = split /\s+/, $line;
        
        # Ensure we have enough fields to parse.
        next unless @fields >= 4;

        # The fields of interest are the 3rd and 4th columns.
        my $local_addr_port = $fields[2];
        my $peer_addr_port = $fields[3];
	#print "lp: $local_addr_port\n";
	#print "pp: $peer_addr_port\n";

        # Use a regular expression to capture the IP address and port number for both local and peer connections.
        my ($local_ip, $local_port) = ($local_addr_port =~ /^(.+):(\d+)$/);
        my ($peer_ip, $peer_port) = ($peer_addr_port =~ /^(.+):(\d+)$/);

        # If both local and peer details were successfully parsed, add them to the connections array.
        if ($local_ip and $local_port and $peer_ip and $peer_port) {
            push @connections, {
                local_ip   => $local_ip,
                local_port => $local_port,
                ip    => $peer_ip,
                port  => $peer_port,
            };
        }
    }
    return @connections;
}

sub latlon_to_xy {
    my ($latitude, $longitude) = @_;
    my $width = $state->{map_width};
    my $height = $state->{map_height};
    my $xOffset = -29.0;
    my $yOffset = 76.0;

    my $x;
    if ($longitude > 90) {
        $x = ($longitude + 180) * ($width / 360) + $xOffset - (($longitude/180) * 15);
    } else {
        $x = ($longitude + 180) * ($width / 360) + $xOffset;
    }

    my $latRad = $latitude * PI / 180;
    my $mercN = log(tan((PI / 4) + ($latRad / 2)));
    my $y = ($height / 2) - ($width * $mercN / (2 * PI)) + $yOffset;
    return ($x, $y);
}

sub get_app_for_ipv4 {
    my ($ipv4_address) = @_;
    return unless $lsof_sudo_required_on;
    my $lsof_command = 'lsof -R -P -i -n';
    open(my $lsof_pipe, "-|", $lsof_command) or die "Can't run lsof: $!";
    while (my $line = <$lsof_pipe>) {
        if ($line =~ /\b\Q$ipv4_address\E\b/) {
            my ($command) = split /\s+/, $line;
            close $lsof_pipe;
            return $command;
        }
    }
    close $lsof_pipe;
    return;
}

sub find_ip_in_geolite_asn {
    my ($geolite_csv_file, $ip_to_check) = @_;
    my $input_ip = Net::IP->new($ip_to_check);
    unless ($input_ip && $input_ip->version() == 4) {
        warn "Invalid IPv4 address provided: $ip_to_check\n";
        return;
    }
    my $input_ip_num = $input_ip->intip();
    my $fh;
    unless (open($fh, '<', $geolite_csv_file)) {
        warn "Could not open GeoLite2 ASN file '$geolite_csv_file': $!\n";
        return;
    }
    my $file_size = -s $geolite_csv_file;
    my ($low, $high) = (0, $file_size);
    my $found_line;
    my $header = <$fh>;
    $low = tell($fh);

    while ($low <= $high) {
        my $mid = int(($low + $high) / 2);
        seek($fh, $mid, SEEK_SET);
        <$fh> if $mid > 0;
        my $line = <$fh>;
        next unless defined $line && $line =~ /\S/;
        my ($network_cidr) = split /,/, $line, 2;
        my $network_ip = Net::IP->new($network_cidr);
        unless ($network_ip) {
            warn "Warning: Could not parse CIDR '$network_cidr' on line: $line";
            $high = $mid - 1;
            next;
        }
        my $start_ip_num = $network_ip->intip();
        my $end_ip_num   = $start_ip_num + ($network_ip->size() - 1);
        if ($input_ip_num >= $start_ip_num && $input_ip_num <= $end_ip_num) {
            $found_line = $line;
            last;
        } elsif ($input_ip_num < $start_ip_num) {
            $high = $mid - 1;
        } else {
            $low = $mid + 1;
        }
    }
    close $fh;
    return unless defined $found_line;
    my @fields = split /,/, $found_line;
    chomp @fields;
    return {
        network         => $fields[0],
        asn             => $fields[1],
        organization    => $fields[2] || "N/A",
    };
}

#======================================================================
# Tk Drawing and Timer Functions
#======================================================================

# This subroutine replaces the 'on_draw' signal handler from Gtk2.
# It manually clears and redraws all items on the canvas.
sub redraw_canvas {
    my ($canvas) = @_;








    # Delete old drawings (dots and text) before drawing new ones.
    $canvas->delete('connection_item');






    my $radius = $state->{dot_diameter} / 2;

    for my $point (@{$state->{current_points}}) {
        my ($x, $y) = latlon_to_xy($point->{lat}, $point->{lon});

        # Determine color based on port (logic is identical to the original)
        my $dot_color = $state->{default_dot_color};
        if ($point->{local_port} == 80 || $point->{local_port} == 443) {
            $dot_color = $state->{http_dot_color};
        } elsif ($point->{port} == 6667 || $point->{port} == 6669 || $point->{port} == 6660) {
            $dot_color = $state->{irc_dot_color};
        } elsif ($point->{port} == 22) {
            $dot_color = $state->{ssh_dot_color};
        } elsif ($point->{port} == 993) {
            $dot_color = $state->{imap_dot_color};
        } elsif ($point->{port} == 31415 || $point->{port} == 31416) {
            $dot_color = $state->{shadsocks_dot_color};
        }

        # 1. Draw the dot using Tk::Canvas->createOval
        $canvas->createOval(
            $x - $radius, $y - $radius, $x + $radius, $y + $radius,
            -fill    => $dot_color,
            -outline => $dot_color, # Use same color for outline to make a solid dot
            -tags    => 'connection_item',
        );

        # 2. Prepare and draw the text label using Tk::Canvas->createText
        my $label = "";
        my $is_http_or_default = 
            ($dot_color eq $state->{http_dot_color}) || 
            ($dot_color eq $state->{default_dot_color});

        if ($is_http_or_default) {
            my $programname = get_app_for_ipv4($point->{ip}) // " ";
            my $asn_info = $state->{ip2asn}->{$point->{ip}};
            unless ($asn_info) {
                my $asn_hash = find_ip_in_geolite_asn($geolite_file, $point->{ip});
                if ($asn_hash) {
                    $asn_info = "$asn_hash->{asn} - $asn_hash->{organization}";
                    $state->{ip2asn}->{$point->{ip}} = $asn_info;
                }
            }
            $asn_info //= "";

            if ($dot_color eq $state->{http_dot_color}) {
                $label = "$programname -> $point->{ip} - $asn_info";
            } else {
                $label = "$programname:$point->{local_port} -> $point->{ip}:$point->{port} - $asn_info";
            }
        }

        my $text_offset_x = $radius + 2;
        $canvas->createText(
            $x + $text_offset_x, $y,
            -text    => $label,
            -anchor  => 'w', # Anchor text to the west (left)
            -fill    => $state->{text_color},
            -font    => '{Sans} 8',
            -tags    => 'connection_item',
        );
    }
}

sub draw_color_key {
    my ($canvas) = @_;

    my @key_items = (
        { label => "My Webserver Hits", color => $state->{http_dot_color} },
        { label => "IRC",               color => $state->{irc_dot_color} },
        { label => "SSH",               color => $state->{ssh_dot_color} },
        { label => "IMAP/S",            color => $state->{imap_dot_color} },
        { label => "Shadowsocks",       color => $state->{shadsocks_dot_color} },
        { label => "Other",             color => $state->{default_dot_color} },
    );

    my $x_start = 15;
    my $y_start = $state->{map_height} - (scalar(@key_items) * 20) - 10;
    my $rect_size = 10;
    my $line_height = 20;
    my $current_y = $y_start;

    for my $item (@key_items) {
        # Draw the color rectangle
        $canvas->createRectangle(
            $x_start, $current_y, $x_start + $rect_size, $current_y + $rect_size,
            -fill => $item->{color},
            -outline => $item->{color},
        );

        # Draw the text label
        $canvas->createText(
            $x_start + $rect_size + 8, $current_y + $rect_size / 2,
            -text => $item->{label},
            -anchor => 'w',
            -fill   => $state->{text_color},
            -font   => '{Sans} 10 bold',
        );

        $current_y += $line_height;
    }
}


# This subroutine replaces the Glib::Timeout callback.
# It fetches new connection data and then triggers a canvas redraw.
sub update_connections_and_redraw {
    my ($canvas) = @_;
    open my $db_fh, '<', $ipv4_db_path or do {
        warn "Could not open IP database '$ipv4_db_path': $!";
        return;
    };
    my $db_size = -s $db_fh;

    my @connections = get_connections();
    my @new_points;

    for my $conn (@connections) {
        my $ip_decimal = ip_to_decimal($conn->{ip});
        next unless $ip_decimal;
        my $location = find_location_on_disk($db_fh, $ip_decimal, $db_size);
        if ($location) {
            push @new_points, {
                ip         => $conn->{ip},
                port       => $conn->{port},
                lat        => $location->{lat},
                lon        => $location->{lon},
                local_ip   => $conn->{local_ip},
                local_port => $conn->{local_port},
            };
        }
    }
    close $db_fh;
    $state->{current_points} = \@new_points;
    
    # Trigger the drawing subroutine
    redraw_canvas($canvas);
}


#======================================================================
# Main Execution Block
#======================================================================

# 1. Create the main window
my $mw = MainWindow->new;
$mw->title("connmap-tk");
$mw->geometry(sprintf("%dx%d", $state->{map_width}, $state->{map_height}));
#$mw->overrideredirect(1); # Creates a borderless, undecorated window
$mw->protocol('WM_DELETE_WINDOW', sub { exit }); # Ensure clean exit

# 2. Create and pack the canvas widget
my $canvas = $mw->Canvas(
    -width  => $state->{map_width},
    -height => $state->{map_height},
    -borderwidth => 0,
    -highlightthickness => 0, # Remove canvas border
)->pack(-fill => 'both', -expand => 1);

# 3. Load and display the background image using Imager
#    This is more robust than using -file directly with Tk::Photo
my $image_obj = Imager->new;
$image_obj->read(file => $image_path)
    or die "Cannot read image '$image_path' with Imager: ", $image_obj->errstr;

# Convert the image to GIF format in memory, as Tk has excellent GIF support
my $gif_data;
$image_obj->write(data => \$gif_data, type => 'gif')
    or die "Cannot convert image to GIF: ", $image_obj->errstr;

# Create the Tk Photo object from the in-memory GIF data
my $image = $mw->Photo(-data => $gif_data);

# Place the image on the canvas.
$canvas->createImage(0, 0, -anchor => 'nw', -image => $image);


# 4. Implement window dragging
# psych! nope.


# 5. Draw the static color key once
draw_color_key($canvas);

# 6. Set up the recurring timer to update connections
# The 'repeat' method is Tk's equivalent of Glib::Timeout->add
$mw->repeat(1000, \&update_connections_and_redraw, $canvas);

# 7. Perform the first update immediately
update_connections_and_redraw($canvas);

# 8. Start the Tk event loop
MainLoop;

exit 0;