diff --git a/file_utils.py b/file_utils.py index 4aea8a3..ec05f11 100644 --- a/file_utils.py +++ b/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.") \ No newline at end of file diff --git a/octopii.py b/octopii.py index b040ed9..16c5c82 100644 --- a/octopii.py +++ b/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) diff --git a/webhook.py b/webhook.py index d60f87e..558f5b0 100644 --- a/webhook.py +++ b/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}')