lib/chronicle/shell/shell_history_extractor.rb in chronicle-shell-0.1.0 vs lib/chronicle/shell/shell_history_extractor.rb in chronicle-shell-0.2.0

@@ -1,24 +1,22 @@
 require 'chronicle/etl'
-require 'tty-command'
 
 module Chronicle
   module Shell
     class ShellHistoryExtractor < Chronicle::ETL::Extractor
       register_connector do |r|
         r.provider = 'shell'
         r.description = 'shell command history'
         r.identifier = 'shell-history'
       end
 
+      setting :filename, default: ENV['HISTFILE']
+      setting :shell, default: 'bash'
+
       BASH_TIMESTAMP_REGEX = /^\#(?<timestamp>[0-9]{10})/
 
       def prepare
-        # TODO: determine if we're working with bash or zsh
-        @shell_name = 'bash'
-
-        @filename = history_filename
         @commands = load_commands
       end
 
       def results_count
         @commands.count
@@ -27,22 +25,31 @@
       def extract
         @commands.each do |command|
           meta = {
             username: username,
             hostname: hostname,
-            shell_name: @shell_name
+            shell_name: @config.shell
           }
           yield Chronicle::ETL::Extraction.new(data: command, meta: meta)
         end
       end
 
       private
 
+      # TODO: modularize the shell-specific stuff
       def history_filename
+        @config.filename || __send__("history_filename_default_#{@config.shell}")
+      end
+
+      def history_filename_default_bash
         File.join(Dir.home, ".bash_history")
       end
 
+      def history_filename_default_zsh
+        File.join(Dir.home, ".zsh_history")
+      end
+
       def username
         @username ||= Etc.getlogin
       end
 
       def hostname
@@ -53,14 +60,14 @@
       end
 
       def load_commands
         commands = []
 
-        loader = "load_commands_from_#{@shell_name}"
-        send(loader) do |command|
-          next if @options[:load_since] && command[:timestamp] < @options[:load_since]
-          next if @options[:load_until] && command[:timestamp] > @options[:load_until]
+        loader = "load_commands_from_#{@config.shell}"
+        __send__(loader) do |command|
+          next if @config.since && command[:timestamp] < @config.since
+          next if @config.until && command[:timestamp] > @config.until
 
           if block_given?
             yield command
           else
             commands << command
@@ -70,11 +77,11 @@
         commands
       end
 
       def load_commands_from_bash
         timestamp = nil
-        File.foreach(@filename) do |line|
+        File.foreach(history_filename) do |line|
           if match = line.scrub.match(BASH_TIMESTAMP_REGEX)
             timestamp = Time.at(match[:timestamp].to_i)
           elsif timestamp
             command = { timestamp: timestamp, command: line.scrub.strip }
             yield command
@@ -82,9 +89,10 @@
         end
       end
 
       def load_commands_from_zsh
         # TODO: implement this
+        raise NotImplementedError
       end
     end
   end
 end