#!/usr/bin/env perl

#   Mosh: the mobile shell
#   Copyright 2012 Keith Winstein
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.

use warnings;
use strict;
use Socket;
use IO::Pty;
use Getopt::Long;

$|=1;

my $client = 'mosh-client';
my $server = 'mosh-server';

my $usage =
qq{Usage: $0 [options] [user@]host
   --client=PATH        mosh client on local machine  (default: "mosh-client")
   --server=PATH        mosh server on remote machine (default: "mosh-server")\n};

GetOptions( 'client=s' => \$client,
	    'server=s' => \$server,
	    'fake-proxy!' => \my $fake_proxy ) or die $usage;

if ( defined $fake_proxy ) {
  use Errno qw(EINTR);
  use IO::Socket::INET;
  use threads;

  my ( $host, $port ) = @ARGV;

  # Resolve hostname
  my $packed_ip = gethostbyname $host;
  if ( not defined $packed_ip ) {
    die "$0: Could not resolve hostname $host\n";
  }
  my $ip = inet_ntoa $packed_ip;

  print STDERR "MOSH IP $ip\n";

  # Act like netcat
  my $sock = IO::Socket::INET->new( PeerAddr => $ip,
				    PeerPort => $port,
				    Proto => "tcp" )
    or die "$0: connect to host $ip port $port: $!\n";

  sub cat {
    my ( $from, $to ) = @_;
    while ( my $n = $from->sysread( my $buf, 4096 ) ) {
      next if ( $n == -1 && $! == EINTR );
      $n >= 0 or die "$0: read: $!\n";
      $to->write( $buf ) or die "$0: write: $!\n";
    }
  }

  my $thr = threads->create(sub { cat $sock, \*STDOUT; $sock->shutdown( 0 ); })
    or die "$0: create thread: $!\n";
  cat \*STDIN, $sock; $sock->shutdown( 1 );
  $thr->join;
  exit;
}

if ( scalar @ARGV != 1 ) {
  die $usage;
}

my $userhost = $ARGV[ 0 ];

# Run SSH and read password
my $pty = new IO::Pty;
my $pty_slave = $pty->slave;

$pty_slave->clone_winsize_from( \*STDIN );

my $pid = fork;
die "$0: fork: $!\n" unless ( defined $pid );
if ( $pid == 0 ) { # child
  close $pty;
  open STDOUT, ">&", $pty_slave->fileno() or die;
  open STDERR, ">&", $pty_slave->fileno() or die;
  close $pty_slave;

  my $s = q{sh -c 'exec "$@" "`set -- $SSH_CONNECTION; echo $3`"' -- } . $server;
  exec 'ssh', '-o', "ProxyCommand=$0 --fake-proxy -- %h %p", '-t', $userhost, '--', $s;
  die "Cannot exec ssh: $!\n";
} else { # server
  my ( $ip, $port, $key );
  $pty->close_slave();
  LINE: while ( <$pty> ) {
    chomp;
    if ( m{^MOSH IP } ) {
      ( $ip ) = m{^MOSH IP (\S+)\s*$} or die "Bad MOSH IP string: $_\n";
    } elsif ( m{^MOSH CONNECT } ) {
      if ( ( $port, $key ) = m{^MOSH CONNECT (\d+?) ([A-Za-z0-9/+]{22})\s*$} ) {
	last LINE;
      } else {
	die "Bad MOSH CONNECT string: $_\n";
      }
    } else {
      print "$_\n";
    }
  }
  waitpid $pid, 0;

  if ( not defined $ip or not defined $port ) {
    die "$0: Did not find mosh server startup message.\n";
  }

  # Now start real mosh client
  $ENV{ 'MOSH_KEY' } = $key;
  exec {$client} ($client, $ip, $port);
}
