#!/usr/bin/perl

use strict;
use warnings;
use IO::Socket::INET;
use Storable qw(nfreeze);
use Time::HiRes qw(sleep);

# --- Configuration ---
my $server_host = '192.168.1.254'; # IP address of the machine running the server script
my $server_port = 6789;        # Port the server is listening on
my $lsof_sudo_required = 1;    # Set to 1 if lsof needs sudo to see all process names

#======================================================================
# Data Gathering Functions
#======================================================================

sub get_connections_with_progs {
    my %conn_to_prog;
    if ($lsof_sudo_required) {
        # MODIFIED: Corrected lsof parsing logic
        my $lsof_cmd = 'sudo lsof -R -P -iTCP -n'; # Using sudo as intended
        open(my $lsof_pipe, "-|", $lsof_cmd) or do {
            warn "lsof command failed: $!. Will proceed without program names.";
            # Return an empty set of connections if lsof fails, or handle differently
            return (); 
        };
        
        # This regex is more robust. It finds lines ending in (ESTABLISHED)
        # and captures the command, peer IP, and peer port directly.
        # ^(\S+)      - Capture the command from the start of the line.
        # .*?         - Non-greedily match characters until the arrow.
        # ->          - Match the literal arrow.
        # ([^:]+)     - Capture the peer IP (anything not a colon).
        # :(\d+)      - Capture the peer port (digits).
        # \s+\(ESTABLISHED\)$ - Match the end of the string.
        while (my $line = <$lsof_pipe>) {
            if ($line =~ /^(\S+).*?->([^:]+):(\d+)\s+\(ESTABLISHED\)$/) {
                my ($command, $peer_ip, $peer_port) = ($1, $2, $3);
                my $key = "$peer_ip:$peer_port";
                $conn_to_prog{$key} = $command;
            }
        }
        close $lsof_pipe;
    }

    # Now use `ss` to get the definitive list of established connections
    my $output = qx(ss -t -n -4 state established);
    my @lines = split /\n/, $output;
    shift @lines; # remove header
    my @connections;
    for my $line (@lines) {
        my @fields = split /\s+/, $line;
        next unless @fields >= 4;
        my ($local_ip, $local_port) = ($fields[2] =~ /^(.+):(\d+)$/);
        my ($peer_ip, $peer_port) = ($fields[3] =~ /^(.+):(\d+)$/);
        if ($local_ip and $local_port and $peer_ip and $peer_port) {
            my $key = "$peer_ip:$peer_port";
            # Look up the program name from the hash we built with lsof
            my $prog = $conn_to_prog{$key} // '';
            
            push @connections, {
                local_ip    => $local_ip,
                local_port  => $local_port,
                ip          => $peer_ip,
                port        => $peer_port,
                programname => $prog, # This will now be populated
            };
        }
    }
    return @connections;
}

#======================================================================
# Main Client Logic
#======================================================================
sub connect_to_server {
    print "Attempting to connect to $server_host:$server_port...\n";
    while (1) {
        my $socket = IO::Socket::INET->new(
            PeerAddr => $server_host,
            PeerPort => $server_port,
            Proto    => 'tcp',
        );
        return $socket if $socket;
        print "Connection failed. Retrying in 5 seconds...\n";
        sleep 5;
    }
}

# --- Main Loop ---
my $socket = connect_to_server();
print "Connected to server. Starting data transmission.\n";

while (1) {
    # 1. Get the current list of connections with program names
    my @connections = get_connections_with_progs();
    
    my $count_with_progs = grep { $_->{programname} ne '' } @connections;
    print "Found " . scalar(@connections) . " established connections. ($count_with_progs with program names)\n";

    # 2. Serialize the data for network transmission
    my $data_packet = nfreeze(\@connections);
    my $len = length($data_packet);
    my $packed_len = pack('N', $len); # 'N' for network byte order

    # 3. Send the data frame (length header + serialized data)
    my $bytes_sent = $socket->syswrite($packed_len . $data_packet);
    
    # 4. Check for disconnection
    unless (defined $bytes_sent && $bytes_sent == (4 + $len)) {
        print "Failed to send data. Server may have disconnected.\n";
        $socket->close();
        $socket = connect_to_server();
        print "Reconnected to server.\n";
        next; # Restart the loop with the new socket
    }
    
    # 5. Wait before sending the next update
    sleep 1;
}