Listssh

bin/singularity

@@ -16,6 +16,32 @@ def print_usage
 END
 end
 
+def prepare_ssh
+  # copy the docker mounted ssh files and chown them to work inside the container
+  if Dir.exists?('/ssh')
+    FileUtils.copy_entry '/ssh', '/root/.ssh'
+    FileUtils.chown_R 'root', 'root', '/root/.ssh'
+
+    # create & populate ssh config file
+    Dir.chdir('/root/.ssh') do
+      identities = (Dir['*'] - ['.', '..']).select do |entry|
+        open(entry).read.include?('BEGIN RSA PRIVATE KEY')
+      end
+
+      File.open('config', 'w') do |config|
+        identities.each do |i|
+          config.write('IdentityFile /root/.ssh/' + i + "\n")
+        end
+        config.write("Host *\n\tServerAliveInterval 60\n\tTCPKeepAlive yes")
+      end
+
+    end
+  else
+    puts "Your ~/.ssh directory did not successfully mount in the docker image that this command runs in."
+    exit
+  end
+end
+
 action = ARGV[0]
 
 unless File.exist?('mesos-deploy.yml') and File.exist?('.mescal.json')
@@ -50,36 +76,18 @@ case action
     Singularity::Runner.new(ARGV, uri).run
 
   when "ssh"
-    # copy the docker mounted ssh files and chown them to work inside the container
-    if Dir.exists?('/ssh')
-      FileUtils.copy_entry '/ssh', '/root/.ssh'
-      FileUtils.chown_R 'root', 'root', '/root/.ssh'
-
-      # create & populate ssh config file
-      Dir.chdir('/root/.ssh') do
-        identities = (Dir['*'] - ['.', '..']).select do |entry|
-          open(entry).read.include?('BEGIN RSA PRIVATE KEY')
-        end
-
-        File.open('config', 'w') do |config|
-          identities.each do |i|
-            config.write('IdentityFile /root/.ssh/' + i + "\n")
-          end
-          config.write("Host *\n\tServerAliveInterval 60\n\tTCPKeepAlive yes")
-        end
-
-      end
-    else
-      puts "Your ~/.ssh directory did not successfully mount in the docker image that this command runs."
-      exit
-    end
+    prepare_ssh
     Singularity::Runner.new(ARGV, uri).run
 
+  when "list-ssh"
+    prepare_ssh
+    Singularity::Request.new(nil, uri, nil).list_ssh
+
   when "help"
     print_usage
 
   else
-    puts "Invalid subcommand '#{action}'"
+    puts "Invalid subcommand \'#{action}\'"
     print_usage
 
 end

lib/singularity/request.rb

@@ -32,16 +32,114 @@ def deploy
       else
         @data['requestId'] = @data['id']
         @data['id'] = "#{@release}.#{Time.now.to_i}"
+        # @data['containerInfo']['docker']['image'] = "#{JSON.parse(File.read('.mescal.json'))['image'].split(':').first}:#{@release}"
         @deploy = {
-         'deploy' => @data,
-         'user' => `whoami`.chomp,
-         'unpauseOnSuccessfulDeploy' => false
+          'deploy' => @data,
+          'user' => `whoami`.chomp,
+          'unpauseOnSuccessfulDeploy' => false
         }
         # deploy the request
         RestClient.post "#{@uri}/api/deploys", @deploy.to_json, :content_type => :json
         puts ' Deploy succeeded: '.green + @data['requestId'].light_blue
       end
     end
 
+    def list_ssh
+      activeTasksList = JSON.parse(RestClient.get "#{@uri}/api/tasks/active", :content_type => :json)
+      mySshTaskList = []
+      taskId = ''
+
+      count = 0
+      activeTasksList.each do |entry|
+        taskId = entry['taskRequest']['request']['id']
+        if taskId.include?("SSH")
+          ip = entry['offer']['url']['address']['ip']
+          port = entry['mesosTask']['container']['docker']['portMappings'][0]['hostPort']
+          mySshTaskList.push(entry)
+          puts "#{count+1}: ".light_green + "#{taskId}: ".light_blue + "root".yellow + " @ ".light_blue + "#{ip}".light_magenta + " : ".light_blue + "#{port}".light_cyan
+          count = count + 1
+        end
+      end
+
+      if count == 0
+        puts "There were no running SSH sessions on #{@uri}"
+        exit 0
+      end
+
+      puts "Would you like to (k)ill or (c)onnect to any of these sessions? (x to exit)"
+      resp = STDIN.gets.chomp
+
+      while !['x','k','kill','c','con','conn','connect'].include?(resp)
+        puts "Incorrect input, please enter c, k, or x"
+        resp = STDIN.gets.chomp
+      end
+
+      case resp
+      when 'x'
+        puts 'Exiting...'.light_magenta
+        exit 0
+      when 'k','kill'
+        puts 'Please enter a comma-separated list of which numbers from the above list you would like to kill (or x to exit)'
+        killList = STDIN.gets.chomp
+        if killList == 'x'
+          exit 0
+        end
+        killList = killList.delete(' ').split(',')
+        killList.each do |task_index|
+          thisTask = mySshTaskList[task_index.to_i-1]
+          puts '!! '.red + 'Are you sure you want to KILL ' + "#{thisTask['taskId']['requestId']}".red + '? (y/n)' + ' !!'.red
+          if STDIN.gets.chomp == 'y'
+            RestClient.delete "#{@uri}/api/requests/request/#{thisTask['taskId']['requestId']}"
+            puts ' KILLED and DELETED: '.red + "#{thisTask['taskId']['requestId']}".light_blue
+          end
+        end
+      when 'c','con','conn','connect'
+        taskIndex = -1
+
+        def getTaskIndex
+          n = 0
+          while n <= 0
+            puts 'Please enter session number to SSH into from the above list (or x to exit)'
+            n = STDIN.gets.chomp
+            n == 'x' ? (puts "Exiting...".light_magenta; exit 0) : (n = Integer(n))
+          end
+          return n-1
+        end
+
+        def pickTask(fromList)
+          taskIndex = getTaskIndex
+
+          puts "SSH into #{fromList[taskIndex]['taskId']['requestId']}? (y = yes, p = pick another task number, or x = exit)"
+          input = STDIN.gets.chomp
+
+          while !["x","y","p"].include?(input)
+            puts "Please enter: y, p, or x"
+            input = STDIN.gets.chomp
+          end
+
+          case input
+          when 'x'
+            puts "Exiting...".light_magenta
+            exit 0
+          when 'p'
+            pickTask(fromList)
+          when 'y'
+            puts "Just a moment... connecting you to the instance."
+          end
+
+        end
+
+        pickTask(mySshTaskList)
+
+        # create fresh Runner, which normally creates a new request when it runs
+        # so we assign values to it to "turn it into" our currently running SSH task
+        runner = Singularity::Runner.new(['ssh'], @uri)
+        runner.thisIsANewRequest = false # so that we don't create & deploy a new request
+        runner.thisTask = mySshTaskList[taskIndex]
+        runner.projectName = mySshTaskList[taskIndex]['taskId']['requestId']
+        # then run it, which only does getIPAndPort and runSsh
+        runner.run
+      end
+    end
   end
 end

lib/singularity/runner.rb

@@ -1,13 +1,14 @@
 module Singularity
   class Runner
-    attr_accessor :request, :ip, :port
+    attr_accessor :request, :ip, :port, :thisIsANewRequest, :thisTask, :projectName
 
     def initialize(commands, uri)
       @commands = commands
       @uri = uri
       @ip = 0
       @port = 0
       @thisTask = ''
+      @thisIsANewRequest = true
 
       mescaljson = JSON.parse(File.read('.mescal.json'))
       @cpus = mescaljson['cpus']
@@ -75,9 +76,12 @@ def initialize(commands, uri)
 
     def run
       exit_code = 0
-      @request.create
-      @request.deploy
-      waitForTaskToShowUp()
+      if @thisIsANewRequest
+        @request.create
+        @request.deploy
+        waitForTaskToShowUp
+      end
+      getIPAndPort
 
       if @commands[0] == 'ssh'
         runSsh
@@ -105,6 +109,9 @@ def waitForTaskToShowUp
           end
         end
       end until @thisTask != ''
+    end
+
+    def getIPAndPort
       @ip = @thisTask['offer']['url']['address']['ip']
       @port = @thisTask['mesosTask']['container']['docker']['portMappings'][0]['hostPort']
     end

lib/singularity/util.rb

@@ -3,6 +3,7 @@
 
 module Singularity
   module Util
+
     def self.port_open?(ip, port, seconds=1)
       Timeout::timeout(seconds) do
         begin
@@ -15,5 +16,6 @@ def self.port_open?(ip, port, seconds=1)
     rescue Timeout::Error
       false
     end
+
   end
 end