#!/usr/bin/env ruby

require 'fileutils'
include FileUtils::Verbose
require 'tins/xt'
include Tins::GO
require 'utils'
require 'pstree'
require 'term/ansicolor'
include Term::ANSIColor
Term::ANSIColor.coloring = STDOUT.tty?

SSH_CONFIG = <<SSH_CONFIG_END
# ~/.ssh/config

Host *
  ForwardX11 = yes
  ControlMaster auto
  ControlPath ~/.ssh/%r@%h:%p.sock
SSH_CONFIG_END

def usage
  puts <<EOT
Usage: #{File.basename($0)} [OPTS] [user@]remote[:port]"

OPTS is one of
  -N                                  list all session names on the specified remote
  -n NAME                             name of the multiplexer session (default: $USER)
  -T                                  list current tunnels
  -e VARIABLE=VALUE [...]             set variables to values
  -t [LHOST:][LPORT:][HOST:PORT|PORT] host:port to tunnel if different from LOCALPORT
  -m screen|tmux                      use sshscreen or tmux as a terminal multiplexer
  -C (ssh|rc)-default|rc              output ssh or rc config file
  -d                                  enable debugging
  -h                                  to display this help

Version is #{File.basename($0)} #{Utils::VERSION}.
EOT
  exit 1
end

def cmd(string)
  $DEBUG and warn "Executing: #{string.inspect}"
  exec string
end

config = Utils::ConfigFile.new

arguments = ARGV
opts = go 't:n:C:m:e:hNTd', arguments

$DEBUG = opts[?d]

if opts[?T]
  tunnels = PSTree.new.select { |pt| pt.user == ENV['USER'] && pt.cmd =~ /ssh.*-L/ }
  tunnels.map! do |pt|
    cmd = pt.cmd
    cmd.gsub!(/(?<=\s)[^\/\s]+@\S+/) { |t| red(t) }
    cmd.gsub!(/-L\s+\S+/) { |t| green(t) }
    "#{yellow(pt.pid.to_s)} #{cmd}"
  end
  STDOUT.puts tunnels
  exit 0
end

case opts[?C]
when 'ssh-default'
  STDOUT.puts SSH_CONFIG; exit
when 'rc-default'
  STDOUT.puts config.to_ruby; exit
when 'rc'
  config.configure_from_paths
  STDOUT.puts config.to_ruby; exit
end

config.configure_from_paths

usage if opts[?h] or arguments.size != 1

if multiplexer = opts[?m]
  config.ssh_tunnel.terminal_multiplexer = multiplexer
end

user_remote = arguments.shift
user, remote, rport =
  case user_remote
  when /\A(?:([^@:]+)@)?([^@:]+)(?::(\d+))?\z/
    user = $1 || ENV['USER']
    user.to_s.empty? and fail "user required to login"
    [ user, $2, $3 || '22' ]
  else
    usage
  end

ssh_dir = File.expand_path('~/.ssh')
mkdir_p ssh_dir
sock_file = "#{ssh_dir}/#{user}@#{remote}:#{rport}.sock"
if env_user = ENV['USER']
  opts[?n] ||= env_user
else
  opts[?n] ||= 'session'
end
if opts[?N]
  cmd "ssh -p #{rport} -S #{sock_file} #{user}@#{remote} #{config.ssh_tunnel.multiplexer_list}"
else
  File.exist? sock_file and rm_f sock_file
  env, tunnels = [], []
  config.ssh_tunnel.copy_paste.full? do |t|
    env << "COPY_REMOTE_HOST_PORT='#{t.bind_address}:#{t.port}'"
    tunnels << "-R #{t}"
  end
  opts[?t].to_a.each do |tunnel_spec|
    if arg = Utils::SshTunnelSpecification.new(tunnel_spec).valid?
      tunnels << "-L #{arg}"
    else
      usage
    end
  end
  config.ssh_tunnel.env.each do |var, val|
    ENV[var.to_s] = val.to_s
  end
  opts[?e].to_a.each do |setting|
    var, val = setting.split('=', 2)
    ENV[var] = val
  end
  cmd "ssh -p #{rport} -Mt  "\
    "-S #{sock_file} #{user}@#{remote} #{tunnels * ' '} "\
      "'env #{env * ' '} #{config.ssh_tunnel.multiplexer_new(opts[?n])} || "\
      "#{config.ssh_tunnel.multiplexer_attach(opts[?n])}'"
end