#!/usr/bin/perl

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

# --- Configuration ---
# Default values. Can be overridden by command-line arguments, ie $ ./connmapperl-client.pl 122.184.120.98 passwordhere
my $server_host = '127.0.0.1'; 
my $server_port = 6789;
my $password    = "passwordhere"; # Default password

# --- NEW: Check for command-line arguments ---
# The first argument is the server host.
if (defined $ARGV[0]) {
    $server_host = $ARGV[0];
    print "Server host set from command line: $server_host\n";
}

# The second argument is the password.
if (defined $ARGV[1]) {
    $password = $ARGV[1];
    print "Password set from command line.\n"; # Avoid printing the actual password for security
}

# --- Reconnection Policy ---
my $initial_retry_delay = 2;     # Start with a 2-second delay
my $max_retry_delay     = 300;   # Cap the delay at 5 minutes (300 seconds)
my $current_retry_delay = $initial_retry_delay;

# --- Handle SIGPIPE ---
# Prevents the script from exiting if the server closes the connection while we write.
$SIG{PIPE} = 'IGNORE';


#======================================================================
# Data Gathering Functions (Unchanged)
#======================================================================
sub get_connections_with_progs {
    my %conn_to_prog;
    my $lsof_cmd = 'lsof -R -P -iTCP -n +c 0';
    open(my $lsof_pipe, "-|", $lsof_cmd) or do {
        warn "lsof command failed: $!. Will proceed without program names.";
    };
    if ($lsof_pipe) {
        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;
    }
    my $output = qx(ss -t -n -4 state established);
    my @lines = split /\n/, $output;
    shift @lines;
    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";
            my $prog = $conn_to_prog{$key} // ''; #/
            push @connections, {
                local_ip => $local_ip, local_port => $local_port,
                ip => $peer_ip, port => $peer_port,
                programname => $prog,
            };
        }
    }
    return @connections;
}

#======================================================================
# Main Client Logic (Unchanged)
#======================================================================

# --- connect_to_server attempts to connect only ONCE ---
sub connect_to_server {
    my $socket = IO::Socket::INET->new(
        PeerAddr => $server_host,
        PeerPort => $server_port,
        Proto    => 'tcp',
        Timeout  => 5, # Set a connection timeout
    );

    unless ($socket) {
        return undef; # Connection failed
    }

    # Connection established, now authenticate
    print "Connection established to $server_host. Authenticating...\n";
    my $pass_ref    = \$password;
    my $auth_packet = nfreeze($pass_ref);
    my $len         = length($auth_packet);
    my $packed_len  = pack('N', $len);

    my $bytes_sent = $socket->syswrite($packed_len . $auth_packet);
    unless (defined $bytes_sent) {
        warn "Failed to send authentication packet. Server may have closed connection.\n";
        $socket->close();
        return undef;
    }
    
    print "Authentication successful.\n";
    return $socket; # Return the fully authenticated socket
}


# --- Main loop contains all reconnection logic ---
my $socket;

while (1) {
    # If the socket isn't defined or connected, enter the reconnection loop.
    unless (defined $socket && $socket->connected) {
        # This inner loop will run until a connection is made.
        while (1) {
            print "Attempting to connect to server at $server_host...\n";
            $socket = connect_to_server();

            if (defined $socket) {
                print "Connection to server successful.\n";
                # On successful connection, reset the retry delay for the future.
                $current_retry_delay = $initial_retry_delay;
                last; # Exit the reconnection loop and proceed to send data.
            }

            # If connection failed, wait and increase the delay for the next attempt.
            print "Connection failed. Retrying in $current_retry_delay seconds...\n";
            sleep $current_retry_delay;
            
            $current_retry_delay *= 2; # Double the delay
            if ($current_retry_delay > $max_retry_delay) {
                $current_retry_delay = $max_retry_delay; # Cap at the max delay
            }
        }
    }

    # --- Data Sending ---
    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";

    my $data_packet = nfreeze(\@connections);
    my $len         = length($data_packet);
    my $packed_len  = pack('N', $len);

    my $bytes_sent = $socket->syswrite($packed_len . $data_packet);
    
    # If syswrite returns undef, the connection is gone.
    unless (defined $bytes_sent) {
        print "Failed to send data. Server has disconnected.\n";
        $socket->close();
        $socket = undef; # Set socket to undef to trigger the reconnection logic
    }
    
    # Wait a second before sending the next update.
    sleep 1;
}