Rotaba/structured-data-anonymizer
0
1 2"""Streamlit app for Presidio + Privy-trained PII models."""3 4import spacy5import en_spacy_pii_distilbert6from spacy_recognizer import CustomSpacyRecognizer7from presidio_analyzer.nlp_engine import NlpEngineProvider8from presidio_anonymizer import AnonymizerEngine9from presidio_analyzer import AnalyzerEngine, RecognizerRegistry10import pandas as pd11from annotated_text import annotated_text12from json import JSONEncoder13import json14import warnings15import streamlit as st16# from streamlit import logger as _logger17import os18import csv19import json20from chatgpt_wrapper import ChatGPT21import time22 23os.environ["TOKENIZERS_PARALLELISM"] = "false"24warnings.filterwarnings('ignore')25 26# from flair_recognizer import FlairRecognizer27 28def load_data(file_location):29 unpacked_string_data = []30 unpacked_url_data = []31 unpacked_json_data = []32 # Read the data back from the CSV file and unpack it33 with open(file_location, mode='r') as csv_file:34 reader = csv.reader(csv_file)35 for row in reader:36 unpacked_string_data.append(row[0])37 unpacked_url_data.append(row[1])38 unpacked_json_data.append(json.loads(row[2]))39 # print("Unpacked string data:", unpacked_string_data)40 # print("Unpacked url data:", unpacked_url_data)41 # print("Unpacked JSON data:", unpacked_json_data)42 return unpacked_string_data, dict(zip(unpacked_string_data, unpacked_json_data))43 44# Helper methods45@st.cache_resource #(allow_output_mutation=True)46def analyzer_engine(use_local=None):47 """Return AnalyzerEngine."""48 49 spacy_recognizer = CustomSpacyRecognizer()50 if use_local:51 # !pip52 # install53 # https: // huggingface.co / beki / en_spacy_pii_distilbert / resolve / main / en_spacy_pii_distilbert - any - py3 - none - any.whl54 55 # Using spacy.load().56 nlp = spacy.load("en_spacy_pii_distilbert")57 58 # Importing as module.59 nlp_engine = en_spacy_pii_distilbert.load()60 else:61 configuration = {62 # print("ENALBEE MODELES")63 "nlp_engine_name": "spacy",64 "models": [65 {"lang_code": "en", "model_name": "en_spacy_pii_distilbert"}],66 }67 68 # Create NLP engine based on configuration69 provider = NlpEngineProvider(nlp_configuration=configuration)70 nlp_engine = provider.create_engine()71 72 registry = RecognizerRegistry()73 # add rule-based recognizers74 registry.load_predefined_recognizers(nlp_engine=nlp_engine)75 registry.add_recognizer(spacy_recognizer)76 # remove the nlp engine we passed, to use custom label mappings77 registry.remove_recognizer("SpacyRecognizer")78 79 analyzer = AnalyzerEngine(nlp_engine=nlp_engine,80 registry=registry, supported_languages=["en"])81 82 # uncomment for flair-based NLP recognizer83 # flair_recognizer = FlairRecognizer()84 # registry.load_predefined_recognizers()85 # registry.add_recognizer(flair_recognizer)86 # analyzer = AnalyzerEngine(registry=registry, supported_languages=["en"])87 return analyzer88 89 90@st.cache_resource#(allow_output_mutation=True)91def anonymizer_engine():92 """Return AnonymizerEngine."""93 return AnonymizerEngine()94 95 96def get_supported_entities():97 """Return supported entities from the Analyzer Engine."""98 return analyzer_engine().get_supported_entities()99 100 101def analyze(**kwargs):102 """Analyze input using Analyzer engine and input arguments (kwargs)."""103 if "entities" not in kwargs or "All" in kwargs["entities"]:104 kwargs["entities"] = None105 return analyzer_engine().analyze(**kwargs)106 107 108def anonymize(text, analyze_results):109 """Anonymize identified input using Presidio Abonymizer."""110 if not text:111 return112 res = anonymizer_engine().anonymize(text, analyze_results)113 return res.text114 115 116def annotate(text, st_analyze_results, st_entities):117 tokens = []118 # sort by start index119 results = sorted(st_analyze_results, key=lambda x: x.start)120 for i, res in enumerate(results):121 if i == 0:122 tokens.append(text[:res.start])123 124 # append entity text and entity type125 tokens.append((text[res.start: res.end], res.entity_type))126 127 # if another entity coming i.e. we're not at the last results element, add text up to next entity128 if i != len(results) - 1:129 tokens.append(text[res.end:results[i+1].start])130 # if no more entities coming, add all remaining text131 else:132 tokens.append(text[res.end:])133 return tokens134 135 136st.set_page_config(page_title="Bitahoy demo", layout="wide")137 138# Side bar -------------------------------------------139# add picture with140st.sidebar.image("structured-data-anonymizer/assets/bitahoy-logo.png", width=200)141 142st.sidebar.markdown(143"""Detect and anonymize PII in structured text such as protocol traces (JSON, SQL, XML etc.)"""144)145 146# dropdown147# titles, json_dict = load_data("structured-data-anonymizer/assets/data_s_short.csv")148# option_list = titles149# option = st.sidebar.selectbox(150# 'Choose an existing structured input?',151# option_list)152 153# dropdown df154# Title,Url,Dict,Prompt,Result155dataframe = pd.read_csv("structured-data-anonymizer/assets/df_data_short.csv")156# select only the third column of the data frame157# select only first column of the data frame158titles = dataframe['Title']159# conver it to a list160titles = titles.values.tolist()161# print(dataframe.iloc[0])162# select first row from dataframe163option_list = titles164# for i in option_list:165# if (dataframe[dataframe['Title'] == i]['Result'].empty):166# i = i + "*"167# print(option_list)168option = st.sidebar.selectbox(169 'Choose an existing structured input?',170 option_list)171 172# # st.sidebar.write('You selected:', option)173# json_dict = dataframe['Dict']174# json_dict = json_dict.values.tolist()175sidebar_text = 'Use small icon-button in right corner to copy input to clipboard'176st.sidebar.write(sidebar_text)177json_dict_option = dataframe[dataframe['Title'] == option]['Dict'].values[0]178 179st.sidebar.code (json_dict_option)180 181#romans complex dropdown182# st.checkbox("Enable/Disable input of existing data", key="disabled")183#184# option = st.selectbox(185# "Choose an existing structured input?",186# option_list,187# # label_visibility=st.session_state.visibility,188# disabled=st.session_state.disabled,189# )190# st.write('You selected:', option)191 192st_entities = st.sidebar.multiselect(193 label="Which entities to look for?",194 options=get_supported_entities(),195 default=['PHONE_NUMBER', 'CREDIT_CARD', 'DATE_TIME', 'MEDICAL_LICENSE', 'US_BANK_NUMBER', 'IP_ADDRESS', 'IBAN_CODE', 'LOCATION', 'EMAIL_ADDRESS']196# default=list(get_supported_entities()),197)198# ['PHONE_NUMBER', 'PERSON', 'CRYPTO', 'AU_TFN', 'ORGANIZATION', 'UK_NHS', 'CREDIT_CARD', 'US_DRIVER_LICENSE',199# 'US_SSN', 'URL', 'AU_MEDICARE', 'DATE_TIME', 'NRP', 'US_PASSPORT', 'MEDICAL_LICENSE', 'US_BANK_NUMBER',200# 'IP_ADDRESS', 'IBAN_CODE', 'US_ITIN', 'AU_ACN', 'SG_NRIC_FIN', 'LOCATION', 'AU_ABN', 'EMAIL_ADDRESS']201# st.sidebar.text(list(get_supported_entities()))202 203st_threshold = st.sidebar.slider(204 label="Acceptance threshold", min_value=0.0, max_value=1.0, value=0.35205)206 207st_return_decision_process = st.sidebar.checkbox(208 "Add analysis explanations in json")209 210api_togg = st.sidebar.checkbox(label='API toggle', value=True)211 212# vertical space213st.sidebar.text("")214# vertical space215st.sidebar.text("")216 217 218st.sidebar.info(219 "Privy is an open source framework for synthetic data generation in protocol trace formats (json, sql, html etc). Presidio is an open source framework for PII detection and anonymization. "220 "For more info visit [privy](https://github.com/pixie-io/pixie/tree/main/src/datagen/pii/privy) and [aka.ms/presidio](https://aka.ms/presidio)"221)222 223 224# Main panel225if 'first_load' not in st.session_state:226 st.session_state['first_load'] = True227 228 229analyzer_load_state = st.info(230 "Starting analyzer and loading model...")231engine = analyzer_engine()232analyzer_load_state.empty()233 234# Initialization235# if 'bot' not in st.session_state:236# st.sidebar.text("init...")237# st.session_state['bot'] = ChatGPT()238# init_prompt = "i'd like you to act like a snobby AI and tell me what you think of my structured data"239# init_answer = st.session_state['bot'].ask(init_prompt)240 241# col?242# Store the initial value of widgets in session state243if "visibility" not in st.session_state:244 st.session_state.visibility = "visible"245 st.session_state.disabled = False246 247col1, col2 = st.columns(2)248 249with col1:250 st.subheader("Input")251 252 sys_name = st.text_area(253 label="Name of the system in question",254 value=option,255 height=1,256 )257 258 st_text = st.text_area(259 label= "Structured text used as input",260 value = """{ "@timestamp":"2022-06-08T16:54:58.849Z", "alienOTX":{ "firewall":{ "action":"Deny", "category":"AlienVaultFirewallNetworkRule", "icmp":{ "request":{ "code":"8" } }, "operation_name":"AzureFirewallNetworkRuleLog", "path": "http://www.example.com/ab001.zip", }, "resource":{ "group":"TEST-FW-RG", "id":"/SUBSCRIPTIONS/23103928-B2CF-472A-8CDB-FR7630006000011234567890189/RESOURCEGROUPS/TEST-FW-RG/PROVIDERS/MICROSOFT.NETWORK/AZUREFIREWALLS/TEST-FW01", "address":"172.24.0.4", "provider":"SonicWall", "number":"040084913373", "sentto": "willh@hotmail.com" }, "subscription_id":"4012888888881881-23103928-B2CF-472A-8CDB-0146E2849129" } }""",261 # value="SELECT shipping FROM users WHERE shipping = '201 Thayer St Providence RI 02912'"262 # "\n\n"263 # "{user: Willie Porter, ip: 192.168.2.80, email: willie@gmail.com}",264 height=300,265 )266 button = st.button("Detect and replace PII")267 st.text("""""")268 269with col2:270 st.subheader("Analyzed results with detected entities highlighted")271 # st.text("Output text with detected entities highlighted")272 with st.spinner("Analyzing..."):273 if button or st.session_state.first_load:274 option = sys_name275 st_analyze_results = analyze(276 text=st_text,277 entities=st_entities,278 language="en",279 score_threshold=st_threshold,280 return_decision_process=st_return_decision_process,281 )282 # """283 # Ugly hack that checks if last 2 chars as Z" and changes the end of the last entity to -1284 # This is done to prevent the anotation to inlcude the quotes for the date4 and breka the json donwtheroad285 # ### TODO: make this less hacky?286 # """287 for i in st_analyze_results:288 # st.write(i)289 # st.write(st_text[i.end - 2:i.end])290 if st_text[i.end-2:i.end] == 'Z"':# and i.type == "DATE_TIME":291 i.end = i.end-1292 continue293 if st_text[i.end-2:i.end] == "Z'":# and i.type == "DATE_TIME":294 i.end = i.end-1295 continue296 # if "'" in st_text[i.start:i.end]:297 # st_analyze_results.remove(i)298 # continue299 # if "," in st_text[i.start:i.end]:300 # st_analyze_results.remove(i)301 # continue302 303 304 annotated_tokens = annotate(st_text, st_analyze_results, st_entities)305 # annotated_tokens306 annotated_text(*annotated_tokens)307 308 # vertical space309 st.text("")310 st.text("")311 312 with st.expander("Show results with replaced PII and detailed results"):313 # st.subheader("Final results with tokens instead if PII")314 # vertical space315 316 if button or st.session_state.first_load:317 st_anonymize_results = anonymize(st_text, st_analyze_results)318 st.write(st_anonymize_results)319 # st.write(st_anonymize_results)320 # try:321 # # st_anonymize_results = ast.literal_eval(st_anonymize_results)322 # st.json(st_anonymize_results) #.replace("'", '"'))323 # except Json Parse Error as e:324 # st.write(st_anonymize_results)325 # vertical space326 st.text("")327 st.subheader("Detailed Findings")328 if st_analyze_results:329 res_dicts = [r.to_dict() for r in st_analyze_results]330 for d in res_dicts:331 d['Value'] = st_text[d['start']:d['end']]332 df = pd.DataFrame.from_records(res_dicts)333 df = df[["entity_type", "Value", "score", "start", "end"]].rename(334 {335 "entity_type": "Entity type",336 "start": "Start",337 "end": "End",338 "score": "Confidence",339 },340 axis=1,341 )342 343 st.dataframe(df, width=1000) # , height=500)344 else:345 st.text("No findings")346 347# st_analyze_results348# end of col349 350# After the columns351 352col5, col6 = st.columns(2)353prompt = "Write a summary for a {} event log, based on the given structured JSON input. Start with an executive summary with a short general description of what is a {}, and then focus on the Key Findings, Monitoring Summary, Incident Summary, Threat Summary and Recommendations. Replace any random " \354 "strings and tokens in angular-brackets with an approximations to make it more human readable: \"{}\" ".format(355 option, option,356 st_anonymize_results)357 358with col5:359 st.subheader("Formatting")360 361 button_create = st.button("Create summary")362 st.markdown(363 "Start with an executive summary and describe what system the log came from, then focus on the Key Findings, Monitoring Summary, Incident Summary, Threat Summary and Recommendations.")364 st.text("""""")365 366 with st.expander("Additional inputs"):367 st_prompt = st.text_area(368 label="Tokenized input with the formatted prompt",369 value=prompt,370 height=200,371 )372 373 write_results = ""374 st_output = st.text_area(375 label="Record results for later use",376 value=write_results,377 height=100,378 )379 button_save = st.button("Save summary to file?")380 st.text("""""")381 382 placeholder_table = st.empty()383 placeholder_table.write("")384 385 init_prompt = """I want you to act as a cyber security analyst expert. I will provide some specific information about concrete incidents, and it will be your job to come up with a coherent summery of the event, described in this log I give you. You can give a short description and then give strategies for protecting this system from malicious actors, based on the incident data I give you. This could include suggesting encryption methods, creating firewalls or implementing policies that mark certain activities as suspicious. Your summery would be used by decision makers to manage the situation, therefore make informed predictions and formulate them precisely in relation to the event I present to you."""386 st_init_prompt = st.text_area(387 label="Initial promopt to focus model",388 value=init_prompt,389 height=100,390 )391 button_reset = st.button("Reset model setup")392 393import random394 395with col6:396 st.subheader("Output incident summary")397 398# effect button_create(button2)399# with st.spinner("button_create..."):400if button_create:401 # load existing promp and results402 if (not api_togg):403 saved_prompt = dataframe[dataframe['Title'] == option]['Prompt'].values[0]404 saved_result = dataframe[dataframe['Title'] == option]['Result'].values[0]405 else:406 saved_prompt = ""407 saved_result = ""408 # check if match to current prompt409 # if re.sub(r"[\n\t\s]*", "", saved_prompt) == re.sub(r"[\n\t\s]*", "", st_prompt):410 md_results = ""411 with col6:412 x = st.empty()413 x.markdown("")414 415 # check if saved_prompt is not of a type float416 if (not isinstance(saved_prompt, float)) and (not api_togg):417 # st.write(saved_prompt)418 with col6:419 with st.spinner('Fetching results...'):420 time.sleep(random.uniform(2.1, 5.8))421 # st.write("Prompt already queried in the past, loading result from database")422 md_results = saved_result423 words = md_results.split()424 num_words = len(words)425 chunk_size = int(random.uniform(2, 6))426 str_placeholder = ""427 428 for i in range(0, num_words, chunk_size):429 chunk = ' '.join(words[i:i + chunk_size])430 str_placeholder = str_placeholder + " " + chunk431 x.markdown(str_placeholder)432 # x.markdown(chunk)433 time.sleep(random.uniform(0.1, 0.6))434 x.markdown(saved_result)435 else:436 # st.write("New prompt, need GPT")437 with col6:438 with st.spinner('Generating, please wait...'):439 bot = ChatGPT()440 # init_answer = bot.ask(init_prompt)441 442 init_points = ""443 for chunk in bot.ask_stream(init_prompt):444 init_points = init_points + "."445 x.markdown(init_points)446 447 x.markdown("")448 449 # st_prompt = "tell me two facts about yourself"450 for chunk in bot.ask_stream(st_prompt):451 md_results = md_results + chunk452 x.markdown(md_results)453 #check if last char of chunk is a new line454 # if "\n" in chunk:455 # x.markdown(md_results)456 # st.markdown(chunk)457 x.markdown(md_results)458 bot._cleanup()459 # md_results = bot.ask(st_prompt) #"Hello, could you tell what is {}?".format(option))460 # print(md_results) # prints the response from chatGPT461 462 # st.write(st_prompt)463 # st.write(saved_prompt)464 # md_results = """No result found""" ##here GPT465 # with col6:466 # # st.subheader("Output incident summary")467 # st.markdown(md_results)468 placeholder_table.write((dataframe.loc[dataframe['Title'] == option]))469 470# if button_reset:471# bot = ChatGPT()472# bot._cleanup()473 474if button_save:475 # dataframe = pd.read_csv("structured-data-anonymizer/assets/df_data_short.csv")476 # save st_prompt and st_output to dataframe in row for Title = json_dict_option477 dataframe.loc[dataframe['Title'] == option, 'Prompt'] = st_prompt478 dataframe.loc[dataframe['Title'] == option, 'Result'] = md_results #st_output479 # st.write(json_dict_option)480 # write dataframe back to the csv file481 dataframe.to_csv("structured-data-anonymizer/assets/df_data_short.csv", index=False)482 st.write("Saved to file")483 st.write(dataframe.loc[dataframe['Title'] == option])484 485# end of document486st.session_state['first_load'] = True487 488 489 490class ToDictListEncoder(JSONEncoder):491 """Encode dict to json."""492 493 def default(self, o):494 """Encode to JSON using to_dict."""495 if o:496 return o.to_dict()497 return []498 499 500if st_return_decision_process:501 st.json(json.dumps(st_analyze_results, cls=ToDictListEncoder))502