##############################################
# UDP Relay
#      March 2004, by Netwallah
# This program can be used as a Client, server, or Relay
# for UDP packets.
# In CLIENT mode, it sends whatever is keyed in to a Dest UDP IP:Port
# In SERVER mode, it receives packets on a port.
# In RELAY mode, it receives packets, and forwards them to a different IP:port.
# UDP only.
# --How to Run  ---
# perl UDPRelay <Mode> <Params>
#  <Mode> is one of   C   S   or  R
#  In Client mode, Params are <Dest IP>  <Dest-port> <Optional-Local-Port>
#  In Server mode, Params are a list of ports to listen on
#  In Relay mode, Params are [<listen-port> <Dest IP> <Dest-port>] ...
#
# It can relay multiple ports with a single instance.
# Relay is bi-directional.
###############################################
use strict;
use warnings;
use IO::Socket;
use IO::Select;
use Data::Dumper;

my $MAX_TO_READ = 8000; #1024;
my $ClientServer = shift or die "Please specify CLIENT SERVER or RELAY mode";
my $verbose =   ($ClientServer =~m/v/i);



my ($datagram,$flags,$dest_port,$dest_addr,$send_socket,
    %Conversation, %GetListenPortFromIP_Port);
$| = 1; # autoflush  STDOUT

if ($ClientServer =~m/^c/i){ #----- Client ----------
  $dest_addr = shift or die "Dest addr required for Client";
  $dest_port = shift or die "dest port required for Client"; # Need extra param
  my $localPort = shift;
  $send_socket=Make_Send_Socket($dest_addr,$dest_port,$localPort);
  Send_Msgs();
}elsif ($ClientServer =~m/^s/i){ # Server -------
   Rcv_Msgs(\&Msg_Handler,1,\@ARGV);
}else{ ## We are a RELAY --- ------------------------
  #Relay comes in triplets - ListenPort, DestAddr, DestPort
  my (@ListenAddrs, @SendSockets, @params);
  @params = ExpandParams(3,@ARGV);
  $verbose and print "RELAY PARAMS: @params\n";
  while (my $listen=shift @params){
    $dest_addr = shift @params or die "Dest addr required for relay";
    $dest_port = shift @params or die "dest port required for RELAY"; # Need extra param
    push @SendSockets, $Conversation{$listen}{FWDSOCKET} = # Where we FWD to
                       Make_Send_Socket($dest_addr,$dest_port);
    push @ListenAddrs,$listen;
    $Conversation{$listen}{SERVIP}   =$dest_addr;
    $Conversation{$listen}{SERVPORT} =$dest_port;
    # Next line allows quick identification of where to send REPLY packets
    $GetListenPortFromIP_Port{$dest_addr . $dest_port} = $listen;
  }
  Rcv_Msgs(\&Relay_Handler,0,\@ListenAddrs,\@SendSockets);
}

###########################################
sub Make_Send_Socket  {
   my ($dest_addr,$dest_port,$local_port) = @_;
   my $send_socket = IO::Socket::INET->new(
                               PeerPort  => $dest_port,
                               PeerAddr  => $dest_addr,
                               LocalPort => $local_port,
                               Proto     => "udp")
      or die "Couldn't be a udp Sender to $dest_addr on dest port $dest_port (Local=$local_port) : $@\n";
   return $send_socket;
}
################################
sub Send_Msgs{
   print "I am the Client - please enter data to send to $dest_addr on port $dest_port:\n";
   $dest_addr or die "Dest IP not specified\n";
   while (<>){
         chomp;
         m/quit/i and exit 0;
         if (m/^len[a-z]* (\d+)\s*$/i){ # Need to generate data of this length..
            my $len = $1 || 1;
            my $txt = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789";
            $_ = substr(   $txt x ((length($txt)+$len)/length($txt) ),   0,  $len);
         };
         $send_socket->send ($_);
         sleep .25; # Wait for response...
         print   "RCV:" . read_socket($send_socket) . "\n";
   }
   exit;
}
#######################################
sub   Rcv_Msgs{
  my $handler_sub = shift;
  my $Echo_Request = shift; # Does he want us to Echo ?
  my @PortList = @{shift()} or die "Please specify one or more ports to listen on";
  my @ExtraHandles = @{shift() || []}; # Existing handles to listen on.
  my $Socket_Selector = IO::Select->new();
  my @readyHandles;
  my %SockHandle_To_Port;
  $verbose and print "Setting up sockets to listen on ports @PortList\n";
  # Setup all receive sockets and selector
  foreach my $port (@PortList){
     my $handle = IO::Socket::INET->new(
                               LocalPort => $port,
                               Proto     => "udp")
        or die "Couldn't be a udp Listener on port $port : $@\n";
    $Socket_Selector->add($handle);
    $SockHandle_To_Port{$handle} = $port;
    # Next line required for RELAY operation...
    $Conversation{$port}{LISTENSOCKET} = $handle;
    $Conversation{$port}{FWDPACKETS}   = 0;
    $Conversation{$port}{REPLYPACKETS} = 0;
  }
     # Throw in the STDIN ...
      ## $Socket_Selector->add(\*STDIN);  # Does not work on Win32
  # We also listen on the extra handles...
  foreach(@ExtraHandles){
       $Socket_Selector->add($_);
       #Hmm - how do we associate these with what we sent ?
  }
  $verbose and print "Going into SERVER mode ==== loop receiving messages:\n";
  #while (my $Sender_Info = $rcv_socket->recv($datagram, $MAX_TO_READ, $flags)) {
 while(@readyHandles = $Socket_Selector->can_read){
    foreach my $h (@readyHandles){
          # Find Port corresponding to handle....
          my $rcvport = $SockHandle_To_Port{$h} || 0;
          my ( $Sender_addr, $Sender_port, $datagram) =  read_socket($h);
          &$handler_sub($Sender_addr, $Sender_port, $datagram,$rcvport);
          $Echo_Request and $h->send("ECHO:$datagram"); # ECHO it
    }
  }
}
###############################
sub read_socket{
    my $h = shift;  # Socket Handle

    my ($datagram,$flags);

    my $Sender_Info = $h->recv($datagram, $MAX_TO_READ, $flags) or return 0;
    my ($Sender_port, $Sender_iaddr_binary) = sockaddr_in($Sender_Info);
    my $Sender_addr = inet_ntoa($Sender_iaddr_binary);
    return $Sender_addr, $Sender_port, $datagram;
}
###############################
sub Msg_Handler{
   my ($peer_addr, $port, $datagram,$rcvport) = @_;
   $verbose and print "Rcv \#$Conversation{$rcvport}{FWDPACKETS} on $rcvport from $peer_addr:$port \t:$datagram\n";

}
##################################
sub Relay_Handler{
   Msg_Handler(@_); # Print it first ..
   my ($sender_addr, $Sender_port, $datagram,$rcvport) = @_;
   my $output_handle;

   if ($rcvport == 0){#This is a response that needs to return to client..
        #Send it to Sender_port's handle
        $rcvport = $GetListenPortFromIP_Port{$sender_addr . $Sender_port}
           or warn  "--Error Relaying Server response to orig client for msg from $sender_addr:$Sender_port\n"
                   . "  :Could not find socket handle corresponding to this IP/Port\n";
       $Conversation{$rcvport}{REPLYPACKETS} ++;
       $verbose and print "--Relaying Server response #$Conversation{$rcvport}->{REPLYPACKETS} to  orig client for msg from " .
                "$sender_addr:$Sender_port back to port $rcvport:$datagram\n";
       $output_handle =  $Conversation{$rcvport}{LISTENSOCKET}; # Normally, a LISTEN socket
   }else{
       $Conversation{$rcvport}{FWDPACKETS} ++;
       $output_handle =  $Conversation{$rcvport}{FWDSOCKET}; # where we Foreward to
   }

   $output_handle->send($datagram); # Relay it
}
###############################################

sub ExpandParams{
  my $pcount = shift or die "How many params do u have"; # How many params at a time
  my @origargs = @_;
  my @flat; # what we return
  my ($maxInstances, $currInstances)= (0,0);

  while ( scalar @origargs){
    my @current = ();
    for (my $p=0; $p < $pcount ; $p++){
       my $currentArg = shift @origargs or die "-- Missing Arg $p of $pcount\n";
       $currInstances = 0;
       my ($start,$fin)= split /-/  , $currentArg;
       $fin = $start unless $fin;
       my $prefix ='';
       $start =~m/^(.+\.)(\d+)$/ and ($prefix,$start,$fin)=($1,$2,substr($fin,length($1)));
       $fin < $start and die "Invalid range: $start - $fin \n";
       #print "Params : $currentArg = $start - $fin\n";
       my $valPos=0;
       for my $val ($start..$fin){
          #push @{$current[$valPos]}, $prefix . $val;
          ${$current[$valPos]}[$p] = $prefix . $val;
       #$verbose and print "     Added at \$current\[$valPos]\[$p] = $prefix . $val" .
       #         "(\$currInstances=$currInstances;\$maxInstances=$maxInstances)\n";
          $valPos++;
          $currInstances ++;
       }
       $maxInstances = $currInstances if $currInstances > $maxInstances;
    }

    #print "Max Inst=$maxInstances\n" , Dumper \@current;

    foreach(@current){
      # Sanity check array sizes ...
      $currInstances = scalar @$_ ; # Number of elements in this isnt
      # Find first an instance where ALL values are available
      my @fullval;
      scalar @$_ == $pcount and @fullval = @$_  for @current;
      #print "$currInstances < $maxInstances at @$_ : FULL VAL=@fullval\n";
      for (my $i=0; $i < $pcount; $i++){
         $$_[$i] = $fullval[$i] unless defined $$_[$i];
      }
       #$verbose and print Dumper \@$_;
       push @flat, @$_;
       #print qq( @{$_} \n);
    }
  }
  return @flat;
}

##################################