#!/usr/bin/perl
use warnings;
use strict;
use POE qw(Wheel::TProxySocketFactory Wheel::ReadWrite Filter::Stream);
use Socket qw(unpack_sockaddr_in inet_ntoa);

# This machine is the man-in-the-middle (called mitm)
# Intercepts connections from unsuspecting hosts (clients) to equally unsuspecting TCP servers (servers)

# Starter session, creates the listening TProxySocketFactory
POE::Session->create(
  inline_states => {
    _start => \&mitm_start,
    _stop => \&mitm_stop,
    _child => \&mitm_child,
    mitm_error => \&mitm_error,
    client_connected => \&client_connected,
  },
  args => [ 1800 ],
);

sub mitm_start {
  my ($heap, $session, $port) = @_[HEAP, SESSION, ARG0];
  $heap->{sessions} = 0;
  $heap->{mitm} = POE::Wheel::TProxySocketFactory->new(
    BindPort => $port,
    Reuse => 'yes',
    SuccessEvent => 'client_connected',
    FailureEvent => 'mitm_error',
  );
  logevent("mitm started on port $port", $session);
}

sub mitm_stop {
  my $session = $_[SESSION];
  logevent("mitm is stopping",$session);
}

sub mitm_child {
  my ($heap, $session, $reason) = @_[HEAP, SESSION, ARG0];
  if ($reason eq 'create') {
    $heap->{sessions}++;
  } elsif ($reason eq 'lose') {
    $heap->{sessions}--;
    logevent("mitm connection closed, " . $heap->{sessions} . " active sessions", $session);
  }
}

sub mitm_error {
  my ($heap, $session, $op, $errnum, $errstr) = @_[HEAP, SESSION, ARG0, ARG1, ARG2];
  logevent("mitm $op error $errnum: $errstr", $session);
  delete $heap->{mitm};
}

sub client_connected {
  my ($heap, $session, $client, $client_addr, $client_port) = @_[HEAP, SESSION, ARG0, ARG1, ARG2];

  # client information is about the source, get the IP address of the source
  my $client_host = inet_ntoa($client_addr);

  # get the address and port of the intended server at the other end of this connection
  my ($server_port, $server_addr) = unpack_sockaddr_in(getsockname($client));
  my $server_host = inet_ntoa($server_addr);
  logevent("mitm got connection from $client_host:$client_port to $server_host, " . $heap->{sessions} . " active sessions", $session);

  # create a new Session to manage this client
  POE::Session->create(
    inline_states => {
      _start => \&setup_client,
      _stop => \&shutdown_client,
      server_connect_fail => \&server_connect_fail,
      server_connected => \&server_connected,
      client_input => \&client_input,
      server_input => \&server_input,
      client_conn_error => \&client_conn_error,
      server_conn_error => \&server_conn_error,
      shutdown_rwwheels => \&shutdown_rwwheels,
    },
    args => [ $client, $client_port, $client_host, $server_port, $server_host ],
  );
}

sub setup_client {
  my ($heap, $session, $client, $client_port, $client_host, $server_port, $server_host) = @_[HEAP, SESSION, ARG0..ARG4];
  $heap->{client} = $client;
  $heap->{client_endpoint} = "$client_host:$client_port";
  $heap->{server_endpoint} = "$server_host:$server_port";
  $heap->{server_wheel} = POE::Wheel::TProxySocketFactory->new(
    RemoteAddress => $server_host,
    RemotePort => $server_port,
    SuccessEvent => 'server_connected',
    FailureEvent => 'server_connect_fail',
  );
}

sub shutdown_client {
  my ($heap, $session) = @_[HEAP, SESSION];
  logevent("connection to " . $heap->{server_endpoint} . " terminated", $session);
}

sub server_connected {
  my ($heap, $session, $server) = @_[HEAP, SESSION, ARG0];
  my $client = delete $heap->{client};

  $heap->{client_wheel} = POE::Wheel::ReadWrite->new(
    Handle => $client,
    Filter => POE::Filter::Stream->new,
    InputEvent => 'client_input',
    ErrorEvent => 'client_conn_error',
  );

  $heap->{server_wheel} = POE::Wheel::ReadWrite->new(
    Handle => $server,
    Filter => POE::Filter::Stream->new,
    InputEvent => 'server_input',
    ErrorEvent => 'server_conn_error',
  );

  logevent("connected to " . $heap->{server_endpoint}, $session);
}

sub server_connect_fail {
  my ($heap, $session, $op, $errnum, $errstr) = @_[HEAP, SESSION, ARG0, ARG1, ARG2];
  logevent("connection to " . $heap->{server_endpoint} . "failed. $op error $errnum: $errstr", $session);
  delete $heap->{client};
  delete $heap->{server_wheel};
}

sub client_input {
  my ($heap, $input) = @_[HEAP, ARG0];
  $heap->{server_wheel}->put($input) if $heap->{server_wheel};
}

sub server_input {
  my ($heap, $input) = @_[HEAP, ARG0];
  $heap->{client_wheel}->put($input) if $heap->{client_wheel};
}

sub client_conn_error {
  my ($kernel, $heap, $session, $op, $errnum, $errstr) = @_[KERNEL, HEAP, SESSION, ARG0, ARG1, ARG2];
  if ($op eq 'read' and $errnum == 0) {
    logevent('client disconnected', $session);
  } else {
    logevent("connection from " . $heap->{client_endpoint} . " failed. $op error $errnum: $errstr", $session);
  }
  $kernel->yield("shutdown_rwwheels");
}

sub server_conn_error {
  my ($kernel, $heap, $session, $op, $errnum, $errstr) = @_[KERNEL, HEAP, SESSION, ARG0, ARG1, ARG2];
  if ($op eq 'read' and $errnum == 0) {
    logevent('server disconnected', $session);
  } else {
    logevent("connection to " . $heap->{server_endpoint} . " failed. $op error $errnum: $errstr", $session);
  }
  $kernel->yield("shutdown_rwwheels");
}

sub shutdown_rwwheels {
  my ($heap) = $_[HEAP];
  my $client = $heap->{client_wheel};
  my $server = $heap->{server_wheel};

  $client->shutdown_input();
  $server->shutdown_input();

  $server->flush if $server->get_driver_out_octets();
  delete $heap->{server_wheel};

  $client->flush() if $client->get_driver_out_octets();
  delete $heap->{client_wheel};
}

sub logmsg { print "[$$] @_ at ", scalar localtime, "\n" }
sub logevent {
  my ($state, $session, $arg) = @_;
  my $id = $session->ID();
  print scalar localtime;
  print " session $id $state ";
  print ": $arg" if (defined $arg);
  print "\n";
}

$poe_kernel->run();
exit 0;