Added options support

Author: rohitcoder

connection.yml.sample

@@ -31,6 +31,9 @@ severity_rules:
   Low:
     - query: "length(matches) <= `5`"
       description: "Detected less than 5 PII or Secrets"
+options:
+  quick_exit: True
+  max_matches: 5 ## Default is 1
 sources:
   redis:
     redis_example:

hawk_scanner/commands/gcs.py

@@ -1,11 +1,7 @@
-import argparse
 from google.cloud import storage
 from rich.console import Console
 from hawk_scanner.internals import system
 import os
-import re
-import time
-import yaml
 
 def connect_google_cloud(args, bucket_name, credentials_file):
     try:
@@ -27,6 +23,13 @@ def execute(args):
     shouldDownload = True
     connections = system.get_connection(args)
 
+    options = connections.get('options', {})
+    quick_exit = options.get('quick_exit', False)
+    max_matches = None
+    if quick_exit:
+        max_matches = options.get('max_matches', 1)
+        system.print_info(args, f"Quick exit enabled with max_matches: {max_matches}")
+
     if 'sources' in connections:
         sources_config = connections['sources']
         gcs_config = sources_config.get('gcs')
@@ -41,6 +44,10 @@ def execute(args):
                     bucket = connect_google_cloud(args, bucket_name, credentials_file)
                     if bucket:
                         for blob in bucket.list_blobs():
+                            if quick_exit and len(results) >= max_matches:
+                                system.print_info(args, f"Quick exit: Found {max_matches} matches, exiting...")
+                                break
+                            
                             file_name = blob.name
                             ## get unique etag or hash of file
                             remote_etag = get_last_update_time(blob)

hawk_scanner/commands/s3.py

@@ -36,6 +36,13 @@ def execute(args):
     shouldDownload = True
     system.print_info(args, f"Running Checks for S3 Sources")
     connections = system.get_connection(args)
+    options = connections.get('options', {})
+    quick_exit = options.get('quick_exit', False)
+    max_matches = None
+    if quick_exit:
+        max_matches = options.get('max_matches', 1)
+        system.print_info(args, f"Quick exit enabled with max_matches: {max_matches}")
+
     if 'sources' in connections:
         sources_config = connections['sources']
         s3_config = sources_config.get('s3')
@@ -52,8 +59,11 @@ def execute(args):
                 if access_key and secret_key and bucket_name:
                     bucket = connect_s3(args, access_key, secret_key, bucket_name)
                     if bucket:
-
                         for obj in bucket.objects.all():
+                            if quick_exit and len(results) >= max_matches:
+                                system.print_info(args, f"Quick exit: Found {max_matches} matches, exiting...")
+                                break
+                            
                             remote_etag = obj.e_tag.replace('"', '')
                             system.print_debug(args, f"Remote etag: {remote_etag}")
                             file_name = obj.key

hawk_scanner/commands/slack.py

@@ -25,6 +25,14 @@ def connect_slack(args, token):
 def check_slack_messages(args, client, patterns, profile_name, channel_types, isExternal, read_from, channel_ids=None, limit_mins=60, archived_channels=False):
     results = []
     try:
+        connection = system.get_connection(args)
+        options = connection.get('options', {})
+        quick_exit = options.get('quick_exit', False)
+        max_matches = None
+        if quick_exit:
+            max_matches = options.get('max_matches', 1)
+            system.print_info(args, f"Quick exit enabled with max_matches: {max_matches}")
+
         team_info = client.team_info()
         workspace_url = team_info["team"]["url"].rstrip('/')
         # Helper function to handle rate limits
@@ -137,6 +145,9 @@ def rate_limit_retry(func, *args, **kwargs):
             messages = rate_limit_retry(client.conversations_history, channel=channel_id, oldest=oldest_time, latest=latest_time)["messages"]
             system.print_debug(args, f"Found {len(messages)} messages in channel {channel_name} ({channel_id})")
             for message in messages:
+                if quick_exit and len(results) >= max_matches:
+                    system.print_info(args, f"Quick exit enabled. Found {max_matches} matches. Exiting...")
+                    return results
                 user = message.get("user", "")
                 text = message.get("text")
                 message_ts = message.get("ts")