Merge c116248 into 43b5847

README.md

@@ -36,6 +36,8 @@ Options:
   [--auto-download], [--no-auto-download]  # Enable or disable auto-download of *.zip file(s)
   [--download-to=DOWNLOAD_TO]              # Directory to download file(s)
                                            # Default: /tmp
+  [--size=N]                               # Number of urlscan.io's results. (Max: 100,000)
+                                           # Default: 100
   [--post-to-slack], [--no-post-to-slack]  # Post a message to Slack if it detects a phishing kit
   [--verbose], [--no-verbose]
                                            # Default: true

lib/miteru/cli.rb

@@ -8,11 +8,12 @@ module Miteru
   class CLI < Thor
     method_option :auto_download, type: :boolean, default: false, desc: "Enable or disable auto-download of *.zip file(s)"
     method_option :download_to, type: :string, default: "/tmp", desc: "Directory to download file(s)"
+    method_option :size, type: :numeric, default: 100, desc: "Number of urlscan.io's results. (Max: 100,000)"
     method_option :post_to_slack, type: :boolean, default: false, desc: "Post a message to Slack if it detects a phishing kit"
     method_option :verbose, type: :boolean, default: true
     desc "execute", "Execute the crawler"
     def execute
-      websites = Crawler.execute(options[:verbose])
+      websites = Crawler.execute(size: options[:size], verbose: options[:verbose])
       websites.each do |website|
         next unless website.has_kit?
 

lib/miteru/crawler.rb

@@ -6,17 +6,23 @@
 module Miteru
   class Crawler
     attr_reader :threads
-    def initialize
+    attr_reader :size
+    attr_reader :verbose
+
+    def initialize(size: 100, verbose: false)
       @threads = 10
+      @size = size
+      @verbose = verbose
+      raise ArgumentError, "size must be less than 100,000" if size > 100_000
     end
 
     def suspicous_urls
-      url = "https://urlscan.io/api/v1/search/?q=certstream-suspicious"
+      url = "https://urlscan.io/api/v1/search/?q=certstream-suspicious&size=#{size}"
       res = JSON.parse(get(url))
       res["results"].map { |result| result.dig("task", "url") }
     end
 
-    def execute(verbose = false)
+    def execute
       pool = Thread.pool(threads)
       websites = []
 
@@ -32,8 +38,8 @@ def execute(verbose = false)
       websites
     end
 
-    def self.execute(verbose = false)
-      new.execute(verbose)
+    def self.execute(size: 100, verbose: false)
+      new(size: size, verbose: verbose).execute
     end
 
     private

spec/cli_spec.rb

@@ -6,6 +6,15 @@
   subject { Miteru::CLI.new }
   before(:each) { ENV.delete "SLACK_WEBHOOK_URL" }
 
+  describe "#execute" do
+    before do
+      allow_any_instance_of(Miteru::Crawler).to receive(:suspicous_urls).and_return([])
+    end
+    it "should not raise any error" do
+      Miteru::CLI.start %w(execute)
+    end
+  end
+
   describe "#download_zip_files" do
     before { WebMock.disable! }
     after { WebMock.enable! }
@@ -14,7 +23,7 @@
       zip_files = ["test.zip"]
 
       expect(Dir.glob("#{base_dir}/*.zip").empty?).to be(true)
-      subject.download_zip_files(url, zip_files, @path)
+      capture(:stdout) { subject.download_zip_files(url, zip_files, @path) }
       expect(Dir.glob("#{base_dir}/*.zip").empty?).to be(false)
     end
   end

spec/crawler_spec.rb

@@ -4,10 +4,26 @@
   include_context "http_server"
   subject { Miteru::Crawler }
   describe "#suspicous_urls" do
-    it "should return an Array" do
-      results = subject.new.suspicous_urls
-      expect(results).to be_an(Array)
-      expect(results.length).to eq(100)
+    context "without 'size' option" do
+      it "should return an Array" do
+        results = subject.new.suspicous_urls
+        expect(results).to be_an(Array)
+        expect(results.length).to eq(100)
+      end
+    end
+    context "with 'size' option" do
+      context "when size <= 100,000" do
+        it "should return an Array" do
+          results = subject.new(size: 200).suspicous_urls
+          expect(results).to be_an(Array)
+          expect(results.length).to eq(200)
+        end
+      end
+      context "when size > 100,000" do
+        it "should raise an ArugmentError" do
+          expect { subject.new(size: 100_001).suspicous_urls }.to raise_error(ArgumentError)
+        end
+      end
     end
   end
   describe "#execute" do