From c56722ffc440da166739254075633f15b61b1ac2 Mon Sep 17 00:00:00 2001 From: umair-rhl Date: Tue, 14 Nov 2023 11:31:58 +0530 Subject: [PATCH 1/3] added the notification feature --- file_utils.py | 3 +++ octopii.py | 12 ++++++++++-- webhook.py | 5 ++++- 3 files changed, 17 insertions(+), 3 deletions(-) diff --git a/file_utils.py b/file_utils.py index 4aea8a3..88b789f 100644 --- a/file_utils.py +++ b/file_utils.py @@ -31,6 +31,7 @@ from bs4 import BeautifulSoup from pdf2image import convert_from_path from PIL import Image +import webhook, octopii def truncate(local_location): characters_per_file = 1232500 @@ -142,6 +143,8 @@ def append_to_output_file(data, file_name): with open(file_name, 'w') as write_file: loaded_json.append(data) + if octopii.notifyURL is not None: + webhook.push_data(json.dumps(loaded_json, indent=4), octopii.notifyURL) write_file.write(json.dumps(loaded_json, indent=4)) except: diff --git a/octopii.py b/octopii.py index af60499..d3be026 100644 --- a/octopii.py +++ b/octopii.py @@ -25,6 +25,7 @@ """ output_file = "output.json" +notifyURL = "" import json, textract, sys, urllib, cv2, os, json, shutil, traceback os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' @@ -116,9 +117,16 @@ def search_pii(file_path): 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() diff --git a/webhook.py b/webhook.py index d60f87e..7b612d7 100644 --- a/webhook.py +++ b/webhook.py @@ -28,7 +28,10 @@ def push_data(data: str, url: str): headers = {'Content-type': 'application/json'} - data = f"{"text":'{data}'}" + if "discord" in url: + data = f"{'content':'{data}'}" + else: + data = f"{'text':'{data}'}" req = requests.post ( url, # Example: https://hooks.slack.com/services/<>/<>/<> From 0fae319593c8faaf7628a83cad10cfa143fcf226 Mon Sep 17 00:00:00 2001 From: 0x4f53 <71916237+0x4f53@users.noreply.github.com> Date: Tue, 14 Nov 2023 09:26:02 +0300 Subject: [PATCH 2/3] made changes to webhook method --- file_utils.py | 7 +++---- octopii.py | 8 +++----- webhook.py | 6 ++---- 3 files changed, 8 insertions(+), 13 deletions(-) diff --git a/file_utils.py b/file_utils.py index 88b789f..a3f8718 100644 --- a/file_utils.py +++ b/file_utils.py @@ -24,14 +24,13 @@ 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 from bs4 import BeautifulSoup from pdf2image import convert_from_path from PIL import Image -import webhook, octopii def truncate(local_location): characters_per_file = 1232500 @@ -143,9 +142,9 @@ def append_to_output_file(data, file_name): with open(file_name, 'w') as write_file: loaded_json.append(data) - if octopii.notifyURL is not None: - webhook.push_data(json.dumps(loaded_json, indent=4), octopii.notifyURL) + if len(octopii.notifyURL) > 0: webhook.push_data(json.dumps(loaded_json, indent=4), octopii.notifyURL) 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 d3be026..4cce595 100644 --- a/octopii.py +++ b/octopii.py @@ -121,7 +121,7 @@ def search_pii(file_path): location = sys.argv[1] # Check for the -notify flag - notify_index = sys.argv.index('-notify') if '-notify' in sys.argv else -1 + 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] @@ -208,11 +208,9 @@ def search_pii(file_path): file_utils.append_to_output_file(results, output_file) 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 7b612d7..a2d31d1 100644 --- a/webhook.py +++ b/webhook.py @@ -28,10 +28,8 @@ def push_data(data: str, url: str): headers = {'Content-type': 'application/json'} - if "discord" in url: - data = f"{'content':'{data}'}" - else: - data = f"{'text':'{data}'}" + if "discord" in url: data = "{\"content\":\"{" + data + "}\"}" + else: data = "{\"text\":\"{" + data + "}\"}" req = requests.post ( url, # Example: https://hooks.slack.com/services/<>/<>/<> From 0bac80b52a075518a4a36e183e671dea41eb3832 Mon Sep 17 00:00:00 2001 From: 0x4f53 <71916237+0x4f53@users.noreply.github.com> Date: Tue, 14 Nov 2023 13:53:34 +0300 Subject: [PATCH 3/3] fixed webhooks --- file_utils.py | 2 -- octopii.py | 10 ++++------ webhook.py | 36 ++++++++++++++++++++++-------------- 3 files changed, 26 insertions(+), 22 deletions(-) diff --git a/file_utils.py b/file_utils.py index a3f8718..ec05f11 100644 --- a/file_utils.py +++ b/file_utils.py @@ -139,10 +139,8 @@ 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) - if len(octopii.notifyURL) > 0: webhook.push_data(json.dumps(loaded_json, indent=4), octopii.notifyURL) write_file.write(json.dumps(loaded_json, indent=4)) except: diff --git a/octopii.py b/octopii.py index 4cce595..799c370 100644 --- a/octopii.py +++ b/octopii.py @@ -30,7 +30,7 @@ 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' @@ -112,7 +112,6 @@ def search_pii(file_path): if __name__ in '__main__': - if len(sys.argv) == 1: print_logo() help_screen() @@ -123,10 +122,8 @@ def search_pii(file_path): # 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 + 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() @@ -206,6 +203,7 @@ def search_pii(file_path): results = search_pii (file_path) print(json.dumps(results, indent=4).replace("{", "").replace("}", "").replace(" ", "")) 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...") diff --git a/webhook.py b/webhook.py index a2d31d1..558f5b0 100644 --- a/webhook.py +++ b/webhook.py @@ -24,19 +24,27 @@ SOFTWARE. """ -import requests - +import requests +import json + def push_data(data: str, url: str): headers = {'Content-type': 'application/json'} - if "discord" in url: data = "{\"content\":\"{" + data + "}\"}" - else: data = "{\"text\":\"{" + data + "}\"}" - - req = requests.post ( - url, # Example: https://hooks.slack.com/services/<>/<>/<> - headers=headers, - data=data, - timeout=7 - ) - - 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) + + 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.raise_for_status() # Raise an HTTPError for bad responses + + print('Scan results sent to webhook.') + except requests.exceptions.RequestException as e: + print(f'Couldn\'t send scan results to webhook. Reason: {e}')