lib/threat_agent/tasks/pwnxy.rb in threat_agent-1.0.0.beta.2 vs lib/threat_agent/tasks/pwnxy.rb in threat_agent-1.0.0.beta.3

@@ -1,51 +1,96 @@
 require 'base64'
 require 'cryptic'
 require 'colorize'
+require 'formatador'
 require 'json'
 require 'thor'
 require 'threat_agent'
 
 module ThreatAgent
   module Tasks
     # A namespace for Pwnxy Thor tasks
     #
     # @author Erran Carey <[e-mail]>
     class Pwnxy < Thor
+      class_option :format, aliases: %w[-f], default: :readable, desc: 'The format to display data in'
       desc 'pwnxy info', 'List information on Pwnxy instances'
       def info
         info = $threat_agent_client.request(:pwnxy_info)
-        # TODO: Add a UI class/method.
-        $stdout.puts info
+        if options[:format].eql? 'json'
+          $stdout.puts info
+        else
+          [:encrypted, :encrypted_iv, :encrypted_key].each { |key| info.delete(key) }
+          Formatador.display_table(info)
+        end
       end
 
-      desc 'pwnxy logs [INSTANCE] [OPTIONS]', 'Show logs for a Pwnxy instance'
-      def logs(identifier = 0)
-        logs = $threat_agent_client.request(:pwnxy_logs, { p: identifier })
-        # TODO: Add a UI class/method.
-        # TODO: Return the logs to the user
+      desc 'pwnxy logs [INSTANCE] [ID] [OPTIONS]', 'Show logs for a Pwnxy instance'
+      method_option :encrypted, aliases: %w[-e], default: false, desc: 'Whether or not to decrypt the logs',type: :boolean
+      def logs(pwnxy_id = 0, id = nil)
+        id = id.to_i
+
+        logs = $threat_agent_client.request(:pwnxy_logs, { p: pwnxy_id })
         if logs.is_a?(Hash) && logs['error']
           $stderr.puts "Threat Agent API Error: #{logs['error']}".red
           exit 255 # This is an API error. Exit with an unspecific code.
         end
 
-        $stdout.puts decrypt(logs)
+        if options[:encrypted]
+          if options[:format].eql? 'json'
+            $stdout.puts (id ? logs[id] : logs).to_json
+          else
+            Formatador.display_table(id ? logs[id] : logs)
+          end
+        else
+          decrypted_logs = decrypt(logs)
+
+          if options[:format].eql? 'json'
+            $stdout.puts (id ? decrypted_logs[id] : decrypted_logs).to_json
+          else
+            # # #
+            # Figure out how to use Formatador or an equiv to make a prettier
+            # table.
+            #
+            # Formatador.display_table(id ? decrypted_logs[id] : decrypted_logs)
+            # # #
+            printable = id ? decrypted_logs[id] : decrypted_logs
+            if printable.is_a? Hash
+              printable.each do |name, value|
+                if value
+                  $stdout.puts "#{name}:", value, '---'
+                end
+              end
+            elsif printable.is_a? Array
+              printable.each do |element|
+                printable.each do |name, value|
+                  if value
+                    $stdout.puts "#{name}:", value, '---'
+                  end
+                end
+              end
+            else
+              $stdout.puts printable.to_s
+            end
+          end
+        end
       end
 
       no_commands do
         def decrypt(logs)
-          keypair = Cryptic::Keypair.new(ThreatAgent::Config[:private_key])
-          private_key = keypair.private_key
+          private_key = OpenSSL::PKey::RSA.new(ThreatAgent::Config[:private_key])
 
           logs.map do |log|
             cipher = OpenSSL::Cipher::Cipher.new('aes-256-cbc')
             cipher.decrypt
             cipher.key = private_key.private_decrypt(Base64.decode64(log['encrypted_key']))
             cipher.iv =  private_key.private_decrypt(Base64.decode64(log['encrypted_iv']))
 
             decrypted_data = cipher.update(Base64.decode64(log['encrypted_data']))
             decrypted_data << cipher.final
-          end.to_json
+
+            JSON.parse(decrypted_data) || {}
+          end
         end
       end
     end
   end
 end