Merge pull request #2 from idoa01/socks

Add SOCKS support

assets/canals.sh

@@ -10,7 +10,7 @@ _canal_complete() {
   # Setup the base level (everything after "canal")
   if [ $COMP_CWORD -eq 1 ]; then
     COMPREPLY=( $(compgen \
-                  -W "adhoc create environment help repo restart session setup start stop update" \
+                  -W "adhoc create environment help repo restart session setup socks start stop update" \
                   -- $cur) )
     return 0
   fi

canals.gemspec

@@ -20,8 +20,8 @@ Gem::Specification.new do |s|
   s.require_path          = "lib"
   s.license               = "MIT"
 
-  s.add_dependency        'thor', '~> 0.19.1'
-  s.add_dependency        'terminal-table', '~> 1.5'
+  s.add_dependency        'thor', '~> 0.20.3'
+  s.add_dependency        'terminal-table', '~> 1.8'
 
   s.add_development_dependency 'rspec', '~> 3.4'
 end

lib/canals/cli.rb

@@ -32,21 +32,21 @@ def __print_version
       method_option :user,         :type => :string, :desc => "The user for the ssh proxy host"
       method_option :bind_address, :type => :string, :desc => "The bind address to connect to"
       def create(name, remote_host, remote_port, local_port=nil)
-        opts = {"name" => name, "remote_host" => remote_host, "remote_port" => remote_port, "local_port" => local_port}.merge(options)
+        opts = {name: name, remote_host: remote_host, remote_port: remote_port, local_port: local_port}.merge(options)
         opts = Canals::CanalOptions.new(opts)
         Canals.create_tunnel(opts)
         say "Tunnel #{name.inspect} created.", :green
       end
 
       desc 'delete NAME', "Delete an existing tunnel; if tunnel is active, stop it first"
       def delete(name)
-        tunnel = Canals.repository.get(name)
+        tunnel = Canals.repository.get(name.to_sym)
         if tunnel.nil?
           say "couldn't find tunnel #{name.inspect}. try using 'create' instead", :red
           return
         end
         tstop(name, silent: true)
-        Canals.repository.delete(name)
+        Canals.repository.delete(name.to_sym)
         say "Tunnel #{name.inspect} deleted.", :green
       end
 
@@ -59,7 +59,7 @@ def delete(name)
       method_option :user,         :type => :string, :desc => "The user for the ssh proxy host"
       method_option :bind_address, :type => :string, :desc => "The bind address to connect to"
       def update(name)
-        tunnel = Canals.repository.get(name)
+        tunnel = Canals.repository.get(name.to_sym)
         if tunnel.nil?
           say "couldn't find tunnel #{name.inspect}. try using 'create' instead", :red
           return
@@ -136,6 +136,21 @@ def adhoc(remote_host, remote_port, local_port=nil)
         tstart(opts)
       end
 
+
+      desc "socks LOCAL_PORT", "Create and run a socks connection"
+      method_option :name,         :type => :string, :desc => "The name to use for the socks tunnel, if not supplied a template will be generated"
+      method_option :env,          :type => :string, :desc => "The proxy environment to use"
+      method_option :hostname,     :type => :string, :desc => "The proxy host we will use to connect through"
+      method_option :user,         :type => :string, :desc => "The user for the ssh socks host"
+      method_option :bind_address, :type => :string, :desc => "The bind address to connect to"
+      def socks(local_port)
+        opts = {"adhoc" => true, "socks" => true, "local_port" => local_port}.merge(options)
+        opts["name"] ||= "__SOCKS__"
+        opts = Canals::CanalOptions.new(opts)
+        opts.name = "SOCKS-adhoc-#{opts.hostname}-#{local_port}" if opts.name == "__SOCKS__"
+        tstart(opts)
+      end
+
       desc "environment SUBCOMMAND", "Environment related command (use 'canal environment help' to find out more)"
       subcommand "environment", Canals::Cli::Environment
 

lib/canals/cli/environment.rb

@@ -3,7 +3,6 @@
 require 'canals/environment'
 require 'thor'
 
-
 module Canals
   module Cli
     class Environment < Thor
@@ -19,8 +18,8 @@ def create(name, hostname)
           host = hostname
           user = nil
         end
-        opts = {"name" => name, "hostname" => host}.merge(options)
-        opts["user"] = user if !user.nil?
+        opts = {name: name, hostname: host}.merge(options)
+        opts[:user] = user if !user.nil?
         env = Canals::Environment.new(opts)
         Canals.repository.add_environment(env)
       end

lib/canals/cli/helpers.rb

@@ -6,11 +6,11 @@ module Canals
   module Cli
     module Helpers
 
-      def tstop(tunnel_opts, silent: false)
+      def tstop(tunnel_opts, remove_from_session: true, silent: false)
         if tunnel_opts.instance_of? String
           tunnel_opts = tunnel_options(tunnel_opts)
         end
-        Canals.stop(tunnel_opts)
+        Canals.stop(tunnel_opts, remove_from_session: remove_from_session)
         say "Tunnel #{tunnel_opts.name.inspect} stopped." unless silent
       end
 

lib/canals/cli/session.rb

@@ -42,13 +42,20 @@ def restart
         end
       end
 
-      desc "stop", "Stop the current session"
+      desc "stop", "Stop the current session (stops and removes from session)"
       def stop
         on_all_canals_in_session(:stop) do |canal|
           tstop(canal)
         end
       end
 
+      desc "suspend", "Suspend the current session (stops and doesn't remove from session)"
+      def suspend
+        on_all_canals_in_session(:suspend) do |canal|
+          tstop(canal, remove_from_session: false)
+        end
+      end
+
       no_commands do
         def on_all_canals_in_session(command, &block)
           return if session_empty?

lib/canals/cli/setup.rb

@@ -19,7 +19,6 @@ def wizard
       desc "completion", "Setup bash completion"
       def completion
         install_completion
-        say "Bash completion script upgraded, use `source #{Canals::Tools::Completion.cmp_file}` to reload it", :red
       end
 
       desc "bind-address", "Setup a global bind address (defaults to 127.0.0.1)"
@@ -32,7 +31,7 @@ def bind_address(bind)
         def setup_first_environment
           say "We'll start by setting up your first environment", :green
           say "An 'environment' is the server you connect your tunnels through. you can have many environments."
-          say "The first environment is the default one used for new connections (but you can always change this default in the future"
+          say "The first environment is the default one used for new connections (but you can always change this default in the future)"
           say ""
           return unless yes? "Wait, should we setup your first environment?", :green
           opts = {}
@@ -82,8 +81,9 @@ def check_install_completion
         end
 
         def install_completion
-          Canals::Tools::Completion.install_completion
-          say "Shell completion installed.", :green
+          first_time = Canals::Tools::Completion.install_completion
+          say "Shell completion installed.", :green if first_time
+          say "Bash completion script #{first_time ? "installed" : "upgraded"}, use `source #{Canals::Tools::Completion.cmp_file}` to reload it", :red
         end
 
         def check(check_result, message)

lib/canals/config.rb

@@ -1,6 +1,7 @@
-require 'psych'
 require 'pathname'
+require 'fileutils'
 require 'forwardable'
+require 'canals/tools/yaml'
 
 module Canals
   class Config
@@ -16,14 +17,12 @@ def initialize(root = nil)
     def load_config(config_file)
       valid_file = config_file && config_file.exist? && !config_file.size.zero?
       return {} if !valid_file
-      return Psych.load_file(config_file)
+      return Canals::Tools::YAML.load_file(config_file)
     end
 
     def save!
       FileUtils.mkdir_p(global_config_file.dirname)
-      File.open(global_config_file, 'w') do |file|
-        file.write(Psych.dump(@config))
-      end
+      Canals::Tools::YAML.dump_file(global_config_file, @config)
     end
 
     private

lib/canals/core.rb

@@ -1,4 +1,5 @@
 require 'open3'
+require 'fileutils'
 
 module Canals
 
@@ -23,7 +24,7 @@ def start(tunnel_opts)
       pid.to_i
     end
 
-    def stop(tunnel_opts)
+    def stop(tunnel_opts, remove_from_session: true)
       if tunnel_opts.instance_of? String
         if (Canals.repository.has?(tunnel_opts))
           tunnel_opts = Canals.repository.get(tunnel_opts)
@@ -32,7 +33,7 @@ def stop(tunnel_opts)
         end
       end
       tunnel_close(tunnel_opts)
-      Canals.session.del(tunnel_opts.name)
+      Canals.session.del(tunnel_opts.name) if remove_from_session
     end
 
     def restart(tunnel_opts)
@@ -59,7 +60,11 @@ def socket_file(tunnel_opts)
 
     def tunnel_start(tunnel_opts)
       FileUtils.mkdir_p("/tmp/canals")
-      cmd = "ssh -M -S #{socket_file(tunnel_opts)} -o 'ExitOnForwardFailure yes' -fnNT -L #{tunnel_opts.bind_address}:#{tunnel_opts.local_port}:#{tunnel_opts.remote_host}:#{tunnel_opts.remote_port} #{tunnel_opts.proxy}"
+      if (tunnel_opts.socks)
+        cmd = "ssh -M -S #{socket_file(tunnel_opts)} -o 'ExitOnForwardFailure=yes' -fnNT -D \"#{tunnel_opts.bind_address}:#{tunnel_opts.local_port}\" #{tunnel_opts.proxy}"
+      else
+        cmd = "ssh -M -S #{socket_file(tunnel_opts)} -o 'ExitOnForwardFailure=yes' -fnNT -L #{tunnel_opts.bind_address}:#{tunnel_opts.local_port}:#{tunnel_opts.remote_host}:#{tunnel_opts.remote_port} #{tunnel_opts.proxy}"
+      end
       system(cmd)
       $?
     end

lib/canals/environment.rb

@@ -1,4 +1,4 @@
-require 'psych'
+require 'canals/tools/yaml'
 
 module Canals
   class CanalEnvironmentError < StandardError; end
@@ -28,7 +28,7 @@ def is_default?
     end
 
     def to_yaml
-      Psych.dump(@args)
+      Canals::Tools::YAML.to_yaml(@args)
     end
 
     def to_hash

lib/canals/options.rb

@@ -1,14 +1,14 @@
-require 'psych'
+require 'canals/tools/yaml'
 
 module Canals
   CanalOptionError = Class.new StandardError
 
   class CanalOptions
     BIND_ADDRESS = "127.0.0.1"
-    attr_reader :name, :remote_host, :remote_port, :local_port, :env_name, :env, :adhoc
+    attr_reader :name, :remote_host, :remote_port, :local_port, :env_name, :env, :adhoc, :socks
 
     # define setters
-    [:name, :local_port, :adhoc].each do |attribute|
+    [:name, :local_port, :adhoc, :socks].each do |attribute|
       define_method :"#{attribute}=" do |value|
         @args[attribute] = value
         self.instance_variable_set(:"@#{attribute}", value)
@@ -23,6 +23,7 @@ def initialize(args)
       @remote_port = @args[:remote_port]
       @local_port = @args[:local_port]
       @adhoc = @args[:adhoc] || false
+      @socks = @args[:socks] || false
       @env_name = @args[:env]
       @env = Canals.repository.environment(@env_name)
     end
@@ -57,7 +58,7 @@ def proxy
     end
 
     def to_yaml
-      Psych.dump(@args)
+      Canals::Tools::YAML.to_yaml(@args)
     end
 
     def to_hash(mode=:basic)
@@ -67,16 +68,21 @@ def to_hash(mode=:basic)
     end
 
     def exploded_options
-      {bind_address: bind_address, hostname: hostname, user: user, pem: pem, proxy: proxy, adhoc: adhoc}
+      {bind_address: bind_address, hostname: hostname, user: user, pem: pem, proxy: proxy, adhoc: adhoc, socks: socks}
     end
 
     private
 
     def validate?(args)
       vargs = args.inject({}){|memo,(k,v)| memo[k.to_sym] = v; memo}
       raise CanalOptionError.new("Missing option: \"name\" in canal creation") if vargs[:name].nil?
-      raise CanalOptionError.new("Missing option: \"remote_host\" in canal creation") if vargs[:remote_host].nil?
-      raise CanalOptionError.new("Missing option: \"remote_port\" in canal creation") if vargs[:remote_port].nil?
+      if vargs[:socks]
+        raise CanalOptionError.new("Missing option: \"local_port\" in canal creation") if vargs[:local_port].nil?
+      else
+        raise CanalOptionError.new("Missing option: \"remote_host\" in canal creation") if vargs[:remote_host].nil?
+        raise CanalOptionError.new("Missing option: \"remote_port\" in canal creation") if vargs[:remote_port].nil?
+      end
+
       vargs[:remote_port] = vargs[:remote_port].to_i
       if vargs[:local_port].nil?
         vargs[:local_port] = vargs[:remote_port]

lib/canals/repository.rb

@@ -1,7 +1,8 @@
-require 'psych'
 require 'pathname'
+require 'fileutils'
 require 'forwardable'
 require 'canals/environment'
+require 'canals/tools/yaml'
 
 module Canals
   class Repository
@@ -27,9 +28,9 @@ def empty?
     end
 
     def add(options, save=true)
-      @repo[TUNNELS][options.name] = options.to_hash
+      @repo[TUNNELS][options.name.to_sym] = options.to_hash
       if options.env_name.nil? && !options.env.nil? && options.env.is_default?
-        @repo[TUNNELS][options.name][:env] = options.env.name
+        @repo[TUNNELS][options.name.to_sym][:env] = options.env.name
       end
       save! if save
     end
@@ -40,37 +41,36 @@ def delete(name, save=true)
     end
 
     def get(name)
+      name = name.to_sym
       return nil if !@repo[:tunnels].has_key? name
       CanalOptions.new(@repo[:tunnels][name])
     end
 
     def has?(name)
-      return @repo[:tunnels].has_key? name
+      return @repo[:tunnels].has_key? name.to_sym
     end
 
     def add_environment(environment, save=true)
       if environment.is_default?
-        @repo[ENVIRONMENTS].each { |name, env| env.delete("default") }
+        @repo[ENVIRONMENTS].each { |name, env| env.delete(:default) }
       end
       if @repo[ENVIRONMENTS].empty?
         environment.default = true
       end
-      @repo[ENVIRONMENTS][environment.name] = environment.to_hash
+      @repo[ENVIRONMENTS][environment.name.to_sym] = environment.to_hash
       save! if save
     end
 
     def save!
       FileUtils.mkdir_p(repo_file.dirname)
-      File.open(repo_file, 'w') do |file|
-        file.write(Psych.dump(@repo))
-      end
+      Canals::Tools::YAML.dump_file(repo_file, @repo)
     end
 
     def environment(name=nil)
       if name.nil?
-        args = @repo[ENVIRONMENTS].select{ |n,e| e["default"] }.values[0]
+        args = @repo[ENVIRONMENTS].select{ |n,e| e[:default] }.values[0]
       else
-        args = @repo[ENVIRONMENTS][name]
+        args = @repo[ENVIRONMENTS][name.to_sym]
       end
       Canals::Environment.new(args) if !args.nil?
     end
@@ -89,7 +89,7 @@ def repo_file
     def load_repository(repository_file)
       valid_file = repository_file && repository_file.exist? && !repository_file.size.zero?
       return { ENVIRONMENTS => {}, TUNNELS => {} } if !valid_file
-      return Psych.load_file(repository_file)
+      return Canals::Tools::YAML.load_file(repository_file)
     end
 
   end