Merge pull request #26 from redhuntlabs/webhook-notification

Webhook notification

file_utils.py

@@ -24,7 +24,7 @@
 SOFTWARE.
 """
 
-import requests, xmltodict, json, requests, cv2, urllib, http, traceback, os, textract
+import requests, xmltodict, json, requests, cv2, urllib, http, traceback, os, textract, webhook, octopii
 from skimage import io
 import numpy as np
 from urllib.request import Request, urlopen, re
@@ -139,10 +139,10 @@ def append_to_output_file(data, file_name):
                 loaded_json = json.loads(read_file.read())
         except: # No file
             print ("\nCreating new file named \'" + file_name + "\' and writing to it.")
-
         with open(file_name, 'w') as write_file:
             loaded_json.append(data)
             write_file.write(json.dumps(loaded_json, indent=4))
 
     except:
+        traceback.print_exc()
         print ("Couldn't write to "+ file_name +". Please check if the path is correct and try again.")

octopii.py

@@ -25,11 +25,12 @@
 """
 
 output_file = "output.json"
+notifyURL = ""
 
 import json, textract, sys, urllib, cv2, os, json, shutil, traceback
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
 from pdf2image import convert_from_path
-import image_utils, file_utils, text_utils
+import image_utils, file_utils, text_utils, webhook
 
 model_file_name = 'models/other_pii_model.h5'
 labels_file_name = 'models/other_pii_model.txt'
@@ -111,14 +112,18 @@ def search_pii(file_path):
 
 
 if __name__ in '__main__':
-
     if len(sys.argv) == 1:
         print_logo()
         help_screen()
         exit(-1)
-
     else:
-        location = sys.argv[1] 
+        location = sys.argv[1]
+
+        # Check for the -notify flag
+        notify_index = sys.argv.index('--notify') if '--notify' in sys.argv else -1
+
+        if notify_index != -1 and notify_index + 1 < len(sys.argv): notifyURL = sys.argv[notify_index + 1]
+        else: notifyURL = None
 
     rules=text_utils.get_regexes()
 
@@ -198,13 +203,12 @@ def search_pii(file_path):
             results = search_pii (file_path)
             print(json.dumps(results, indent=4))
             file_utils.append_to_output_file(results, output_file)
+            if notifyURL != None: webhook.push_data(json.dumps(results), notifyURL)
             print ("\nOutput saved in " + output_file)
 
-        except textract.exceptions.MissingFileError:
-            print ("\nCouldn't find file '" + file_path + "', skipping...")
+        except textract.exceptions.MissingFileError: print ("\nCouldn't find file '" + file_path + "', skipping...")
 
-        except textract.exceptions.ShellError:
-            print ("\nFile '" + file_path + "' is empty or corrupt, skipping...")
+        except textract.exceptions.ShellError: print ("\nFile '" + file_path + "' is empty or corrupt, skipping...")
 
     if temp_exists: shutil.rmtree(temp_dir)
 

webhook.py

@@ -24,18 +24,27 @@
 SOFTWARE.
 """
 
-import requests    
-
+import requests
+import json
+
 def push_data(data: str, url: str):
     headers = {'Content-type': 'application/json'}
-    data = f"{"text":'{data}'}"
+
+    if "discord" in url:
+        payload = {"content": data}
+    else:
+        payload = {"text": data}
+
+    try:
+        req = requests.post(
+            url,
+            headers=headers,
+            json=payload,  # Use the json parameter to send JSON data
+            timeout=7
+        )
 
-    req = requests.post (
-        url, # Example: https://hooks.slack.com/services/<>/<>/<>
-        headers=headers, 
-        data=data, 
-        timeout=7
-    )
+        req.raise_for_status()  # Raise an HTTPError for bad responses
 
-    if req is not None and req.status_code == 200: print('Scan results sent to webhook.')
-    else: print('Couldn\'t send scan results to webhook. Reason: ' + req.text)
+        print('Scan results sent to webhook.')
+    except requests.exceptions.RequestException as e:
+        print(f'Couldn\'t send scan results to webhook. Reason: {e}')