##############################################
# 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> and  <Dest-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.
###############################################
use strict;
use warnings;
use IO::Socket;
use IO::Select;

my $MAX_TO_READ = 1024;
my $ClientServer = shift or die "Please specify CLIENT SERVER or RELAY mode";



my ($datagram,$flags,$dest_port,$dest_addr,$send_socket,
    %Conversation);
$| = 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
  $send_socket=Make_Send_Socket($dest_addr,$dest_port);
  Send_Msgs(); # Send_and_Rcv_Forked();
}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);;
  while (my $listen=shift){
    $dest_addr = shift or die "Dest addr required for relay";
    $dest_port = shift or die "dest port required for RELAY"; # Need extra param
    push @SendSockets, #  $Relay{$listen}{SENDSOCKET} =
                       $Conversation{$listen}{SERVSOCKET} =
                       Make_Send_Socket($dest_addr,$dest_port);
    push @ListenAddrs,$listen;
    $Conversation{$listen}{SERVIP}   =$dest_addr;
    $Conversation{$listen}{SERVPORT} =$dest_port;
  }
  Rcv_Msgs(\&Relay_Handler,0,\@ListenAddrs,\@SendSockets);
}

###########################################
sub Make_Send_Socket  {
   my ($dest_addr,$dest_port) = @_;
   my $send_socket = IO::Socket::INET->new(
                               PeerPort      => $dest_port,
                               PeerAddr      => $dest_addr,
                               Proto     => "udp")
      or die "Couldn't be a udp Sender to $dest_addr on port $dest_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;
         #send SOCKET,MSG,FLAGS,TO
         length($_) and $send_socket->send ($_);
         sleep .2; # Wait for response...
         #if (IO::Select->select ( [$send_socket], undef, undef , 2 )){ # check for response

            print   "RCV:" . read_socket($send_socket, $datagram) . "\n";
         #}
   }
   exit;
}
#######################################
sub Send_and_Rcv_Forked{    # Broken
   print "I am the Forked Client - please enter data to send to $dest_addr on port $dest_port:\n";
   $dest_addr or die "Dest IP not specified\n";
   defined(my $child = fork) or  die "fork() failed: $!";
   if ($child == 0){
       #Child process
       while (my $data=read_socket($send_socket, $datagram)){
          print "RCV $data\n";
       }
       exit;
   }else{
      print "Parent here ..Child=$child \n";
        while (<>){
         chomp;
         m/quit/i and exit 0;
         #send SOCKET,MSG,FLAGS,TO
         $send_socket->send ($_);
         }
    };
}
#######################################
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;
  # 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{$Conversation{$port}{SERVPORT}}{CLIENTSOCKET} = $handle;
  }
     # 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 ?
  }
  print "I am the SERVER ====Go into a 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 Process_Echo_And_commands($Echo_Request,$h,$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) = @_;
   print "Rcv $rcvport from $peer_addr:$port \t:$datagram\n";
}
##################################
sub  Process_Echo_And_commands{
   my ($Options,$sock,$Msg) = @_;

   $sock->send("ECHO:$Msg");
     #Handle any recognized commands..
   if ($Msg =~m/^CMD\s+(.+)/){
      my $result = `$1`;
      print "\t$1:$result\n";
      $sock->send($result)
   }
}
##################################
sub Relay_Handler{
   Msg_Handler(@_); # Print it first .. No parsing
   my ($sender_addr, $Sender_port, $datagram,$rcvport) = @_;
   #my $output_handle =  $Relay{$rcvport}{SENDSOCKET};
   my $output_handle =  $Conversation{$rcvport}{SERVSOCKET};
   if ($rcvport == 0){#This is a response that needs to return to client..
        #Send it to Sender_port's handle
        #print "--Relaying Server response to  orig client for $Sender_port:$datagram\n";
        $output_handle = $Conversation{$Sender_port}{CLIENTSOCKET};
   }
   $output_handle->send($datagram); # Relay it
}