#!/usr/bin/env ruby

SHAFT_CONFIG = File.join(Dir.home, ".shaft")
SHAFT_ACTIVE = File.join(Dir.home, ".shaft.active")

require 'rubygems'
require 'thor'
require 'yaml'

class ShaftCLI < Thor
  include Thor::Actions

  map "-h" => :help
  map "-l" => :active

  desc "active", "Lists active tunnels"
  method_options %w( short -s ) => :boolean
  def active
    active = get_active
    unless active.empty?
      if options[:short]
        say active.keys.join(" ")
      else
        say "Listing currently active tunnels:"
        print_table active
      end
    else
      say "No tunnels are currently active." unless options[:short]
    end
  end

  desc "all", "Lists all available tunnels"
  method_options %w( short -s ) => :boolean
  def all
    tunnels = get_config.keys
    if options[:short]
      say tunnels.join(' ')
    else
      say "Listing all available tunnels:"
      print_in_columns tunnels
    end
  end

  desc "completions", "Prints setup code for Zsh completions"
  def completions
    fpath = File.expand_path File.join(File.dirname(__FILE__), '..', 'completions')
    puts "fpath=(#{fpath} $fpath)"
    puts "compinit"
  end

  desc "start [NAME]", "Starts a tunnel"
  method_options :name => :string
  def start(name)
    active = get_active
    if active.has_key? name
      say "Error: tunnel '#{name}' already active!"
    else
      c = get_tunnel(name)
      unless c.nil?
        begin
          port = c['port']
          bind = "#{c['bind']['client-port']}:#{c['bind']['host']}:#{c['bind']['host-port']}"
          host = "#{c['username']}@#{c['host']}"
        rescue NoMethodError
          error "Tunnel configuration for '#{name}' appears to be invalid!"
          return
        end

        used = local_port_used?(c['bind']['client-port'])
        if used
          error "Local port #{port} is used by:\n#{used}"
          return
        end

        say "Starting tunnel '#{name}'..."
        pid = Process.spawn("ssh -N -p #{port} #{host} -L #{bind}")
        Process.detach pid
        say "Started with pid #{pid}."

        active[name] = pid
        set_active(active)
      else
        error "Tunnel '#{name}' not found!"
      end
    end
  end

  desc "stop [NAME]", "Stops a tunnel"
  method_options :name => :string
  def stop(name)
    active = get_active
    if active.has_key?(name)
      say "Stopping tunnel '#{name}' at pid #{active[name]}..."
      begin
        Process.kill "INT", active[name]
      rescue Errno::ESRCH
        say "Tunnel wasn't active (zombie shaft item)."
      end

      # verify killing
      tunnel = get_tunnel(name)
      if local_port_used?(tunnel['bind']['client-port'])
        error "Could not stop tunnel process!"
      else
        say "Stopped."

        # set as inactive
        active.delete(name)
        set_active(active)
      end
    else
      error "Tunnel '#{name}' does not seem to be active!"
    end
  end

  desc "restart [NAME]", "Restarts a tunnel"
  method_options :name => :string
  def restart(name)
    stop(name) && start(name)
  end

  private
  def get_config
    @config ||= read_yaml(SHAFT_CONFIG)
  end

  def get_active
    @active ||= read_yaml(SHAFT_ACTIVE)
  end

  def get_tunnel(name)
    get_config[name] || nil
  end

  def set_active(active)
    File.open(SHAFT_ACTIVE, 'w') { |out|
      YAML.dump(active, out)
    }
  end

  def read_yaml(file)
    if File.directory?(SHAFT_CONFIG)
      error "Shaft v0.7 and up uses a single-file config.\nConsult the Readme at http://git.io/Zs3viQ ."
      {}
    elsif File.exists?(file)
      YAML::load(File.open(file))
    else
      {}
    end
  end

  def local_port_used?(port)
    usage = `lsof -n -i4TCP:#{port} | grep LISTEN`
    if usage.length > 0
      usage
    else
      false
    end
  end

end

ShaftCLI.start