Merge pull request #4 from ninoseki/feat-ordering-option

feat: add order-by option

lib/ukemi.rb

@@ -28,4 +28,6 @@ def services
 
 require "ukemi/moderator"
 
+require "ukemi/configuration"
+
 require "ukemi/cli"

lib/ukemi/cli.rb

@@ -6,8 +6,11 @@
 module Ukemi
   class CLI < Thor
     desc "lookup [IP|DOMAIN]", "Lookup passive DNS services"
+    method_option :order_by, type: :string, desc: "Ordering of the passve DNS resolutions (last_seen or first_seen)", default: "-last_seen"
     def lookup(data)
       data = refang(data)
+      set_ordering options["order_by"]
+
       result = Moderator.lookup(data)
       puts JSON.pretty_generate(result)
     end
@@ -16,6 +19,17 @@ def lookup(data)
       def refang(data)
         data.gsub("[.]", ".").gsub("(.)", ".")
       end
+
+      def set_ordering(order_by)
+        parts = order_by.split("-")
+        ordering_key = parts.last
+        sort_order = parts.length == 2 ? "DESC" : "ASC"
+
+        Ukemi.configure do |config|
+          config.ordering_key = ordering_key
+          config.sort_order = sort_order
+        end
+      end
     end
   end
 end

lib/ukemi/configuration.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Ukemi
+  class Configuration
+    attr_accessor :ordering_key
+    attr_accessor :sort_order
+
+    def initialize
+      @ordering_key = "last_seen"
+      @sort_order = "DESC"
+    end
+  end
+
+  class << self
+    def configuration
+      @configuration ||= Configuration.new
+    end
+
+    attr_writer :configuration
+
+    def configure
+      yield configuration
+    end
+  end
+end

lib/ukemi/moderator.rb

@@ -44,10 +44,17 @@ def format(records)
           }
         ]
       end.to_h
-      # Sort by last_seen (desc)
+
+      # Sorting
+      ordering_key = Ukemi.configuration.ordering_key.to_sym
+      sort_order = Ukemi.configuration.sort_order
       formatted.sort_by do |_key, hash|
-        last_seen = hash.dig(:last_seen)
-        last_seen ? -convert_to_unixtime(last_seen) : -1
+        value = hash.dig(ordering_key)
+        if sort_order == "DESC"
+          value ? -convert_to_unixtime(value) : -1
+        else
+          value ? convert_to_unixtime(value) : Float::MAX.to_i
+        end
       end.to_h
     end
 

spec/cli_spec.rb

@@ -33,4 +33,27 @@
       expect(subject.refang("1.1.1(.)1")).to eq("1.1.1.1")
     end
   end
+
+  describe "#set_ordering" do
+    subject { described_class.new }
+
+    after do
+      Ukemi.configure do |config|
+        config.ordering_key = "last_seen"
+        config.sort_order = "DESC"
+      end
+    end
+
+    it do
+      subject.set_ordering("-first_seen")
+      expect(Ukemi.configuration.ordering_key).to eq("first_seen")
+      expect(Ukemi.configuration.sort_order).to eq("DESC")
+    end
+
+    it do
+      subject.set_ordering("first_seen")
+      expect(Ukemi.configuration.ordering_key).to eq("first_seen")
+      expect(Ukemi.configuration.sort_order).to eq("ASC")
+    end
+  end
 end

spec/moderator_spec.rb

@@ -28,7 +28,7 @@
     allow(Ukemi::Services::SecurityTrails).to receive(:new).and_return(st_mock)
 
     allow(circl_mock).to receive(:lookup).and_return(
-      [Ukemi::Record.new(data: "3.3.3.3", last_seen: "2017-01-01", source: "CIRCL")]
+      [Ukemi::Record.new(data: "3.3.3.3", first_seen: "2017-01-01", last_seen: "2017-01-01", source: "CIRCL")]
     )
     allow(circl_mock).to receive(:configurated?).and_return(true)
     allow(Ukemi::Services::CIRCL).to receive(:new).and_return(circl_mock)
@@ -55,5 +55,41 @@
       last_seen = results.dig(first, :last_seen)
       expect(last_seen).to eq("2019-01-01")
     end
+
+    context "when given an ordering options" do
+      after do
+        Ukemi.configure do |config|
+          config.ordering_key = "last_seen"
+          config.sort_order = "DESC"
+        end
+      end
+
+      it do
+        Ukemi.configure do |config|
+          config.ordering_key = "last_seen"
+          config.sort_order = "ASC"
+        end
+        results = subject.lookup("example.com")
+        expect(results.keys.first).to eq("3.3.3.3")
+      end
+
+      it do
+        Ukemi.configure do |config|
+          config.ordering_key = "first_seen"
+          config.sort_order = "DESC"
+        end
+        results = subject.lookup("example.com")
+        expect(results.keys.first).to eq("2.2.2.2")
+      end
+
+      it do
+        Ukemi.configure do |config|
+          config.ordering_key = "first_seen"
+          config.sort_order = "ASC"
+        end
+        results = subject.lookup("example.com")
+        expect(results.keys.first).to eq("3.3.3.3")
+      end
+    end
   end
 end