added app cahnges and new dataframe'

app.py

@@ -2,6 +2,7 @@
 """Streamlit app for Presidio + Privy-trained PII models."""
 
 import spacy
+import en_spacy_pii_distilbert
 from spacy_recognizer import CustomSpacyRecognizer
 from presidio_analyzer.nlp_engine import NlpEngineProvider
 from presidio_anonymizer import AnonymizerEngine
@@ -12,10 +13,16 @@ from json import JSONEncoder
 import json
 import warnings
 import streamlit as st
+# from streamlit import logger as _logger
 import os
 import csv
+import json
+from chatgpt_wrapper import ChatGPT
+import time
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 warnings.filterwarnings('ignore')
+
 # from flair_recognizer import FlairRecognizer
 
 def load_data(file_location):
@@ -35,22 +42,32 @@ def load_data(file_location):
     return unpacked_string_data, dict(zip(unpacked_string_data, unpacked_json_data))
 
 # Helper methods
-@st.cache(allow_output_mutation=True)
-def analyzer_engine():
+@st.cache_resource #(allow_output_mutation=True)
+def analyzer_engine(use_local=None):
     """Return AnalyzerEngine."""
 
     spacy_recognizer = CustomSpacyRecognizer()
-
-    configuration = {
-        # print("ENALBEE MODELES")
-        "nlp_engine_name": "spacy",
-        "models": [
-            {"lang_code": "en", "model_name": "en_spacy_pii_distilbert"}],
-    }
-
-    # Create NLP engine based on configuration
-    provider = NlpEngineProvider(nlp_configuration=configuration)
-    nlp_engine = provider.create_engine()
+    if use_local:
+        # !pip
+        # install
+        # https: // huggingface.co / beki / en_spacy_pii_distilbert / resolve / main / en_spacy_pii_distilbert - any - py3 - none - any.whl
+
+        # Using spacy.load().
+        nlp = spacy.load("en_spacy_pii_distilbert")
+
+        # Importing as module.
+        nlp_engine = en_spacy_pii_distilbert.load()
+    else:
+        configuration = {
+            # print("ENALBEE MODELES")
+            "nlp_engine_name": "spacy",
+            "models": [
+                {"lang_code": "en", "model_name": "en_spacy_pii_distilbert"}],
+        }
+
+        # Create NLP engine based on configuration
+        provider = NlpEngineProvider(nlp_configuration=configuration)
+        nlp_engine = provider.create_engine()
 
     registry = RecognizerRegistry()
     # add rule-based recognizers
@@ -70,7 +87,7 @@ def analyzer_engine():
     return analyzer
 
 
-@st.cache(allow_output_mutation=True)
+@st.cache_resource#(allow_output_mutation=True)
 def anonymizer_engine():
     """Return AnonymizerEngine."""
     return AnonymizerEngine()
@@ -119,24 +136,48 @@ def annotate(text, st_analyze_results, st_entities):
 st.set_page_config(page_title="Bitahoy demo", layout="wide")
 
 # Side bar -------------------------------------------
+# add picture with
+st.sidebar.image("structured-data-anonymizer/assets/bitahoy-logo.png", width=200)
+
 st.sidebar.markdown(
-"""
-Detect and anonymize PII in structured text such as protocol traces (JSON, SQL, XML etc.)
-"""
+"""Detect and anonymize PII in structured text such as protocol traces (JSON, SQL, XML etc.)"""
 )
-# add picture with
-st.sidebar.image("assets/bitahoy-logo.png", width=200)
 
 # dropdown
-titles, json_dict = load_data("assets/data_sorted.csv")
+# titles, json_dict = load_data("structured-data-anonymizer/assets/data_s_short.csv")
+# option_list = titles
+# option = st.sidebar.selectbox(
+#     'Choose an existing structured input?',
+#     option_list)
+
+# dropdown df
+# Title,Url,Dict,Prompt,Result
+dataframe = pd.read_csv("structured-data-anonymizer/assets/df_data_short.csv")
+# select only the third column of the data frame
+# select only first column of the data frame
+titles = dataframe['Title']
+# conver it to a list
+titles = titles.values.tolist()
+# print(dataframe.iloc[0])
+# select first row from dataframe
 option_list = titles
+# for i in option_list:
+#     if (dataframe[dataframe['Title'] == i]['Result'].empty):
+#         i = i + "*"
+# print(option_list)
 option = st.sidebar.selectbox(
     'Choose an existing structured input?',
     option_list)
 
-# st.sidebar.write('You selected:', option)
-st.sidebar.code (json_dict[option])
-st.sidebar.write('Use button to copy input to clipboard')
+# # st.sidebar.write('You selected:', option)
+# json_dict = dataframe['Dict']
+# json_dict = json_dict.values.tolist()
+sidebar_text = 'Use small icon-button in right corner to copy input to clipboard'
+st.sidebar.write(sidebar_text)
+json_dict_option = dataframe[dataframe['Title'] == option]['Dict'].values[0]
+
+st.sidebar.code (json_dict_option)
+
 #romans complex dropdown
 # st.checkbox("Enable/Disable input of existing data", key="disabled")
 #
@@ -151,8 +192,13 @@ st.sidebar.write('Use button to copy input to clipboard')
 st_entities = st.sidebar.multiselect(
     label="Which entities to look for?",
     options=get_supported_entities(),
-    default=list(get_supported_entities()),
+    default=['PHONE_NUMBER', 'CREDIT_CARD', 'DATE_TIME', 'MEDICAL_LICENSE', 'US_BANK_NUMBER', 'IP_ADDRESS', 'IBAN_CODE', 'LOCATION', 'EMAIL_ADDRESS']
+#     default=list(get_supported_entities()),
 )
+# ['PHONE_NUMBER', 'PERSON', 'CRYPTO', 'AU_TFN', 'ORGANIZATION', 'UK_NHS', 'CREDIT_CARD', 'US_DRIVER_LICENSE',
+# 'US_SSN', 'URL', 'AU_MEDICARE', 'DATE_TIME', 'NRP', 'US_PASSPORT', 'MEDICAL_LICENSE', 'US_BANK_NUMBER',
+# 'IP_ADDRESS', 'IBAN_CODE', 'US_ITIN', 'AU_ACN', 'SG_NRIC_FIN', 'LOCATION', 'AU_ABN', 'EMAIL_ADDRESS']
+# st.sidebar.text(list(get_supported_entities()))
 
 st_threshold = st.sidebar.slider(
     label="Acceptance threshold", min_value=0.0, max_value=1.0, value=0.35
@@ -161,7 +207,7 @@ st_threshold = st.sidebar.slider(
 st_return_decision_process = st.sidebar.checkbox(
     "Add analysis explanations in json")
 
-button = st.sidebar.button("Detect PII")
+api_togg = st.sidebar.checkbox(label='API toggle', value=True)
 
 # vertical space
 st.sidebar.text("")
@@ -176,11 +222,22 @@ st.sidebar.info(
 
 
 # Main panel
+if 'first_load' not in st.session_state:
+    st.session_state['first_load'] = True
+
+
 analyzer_load_state = st.info(
     "Starting analyzer and loading model...")
 engine = analyzer_engine()
 analyzer_load_state.empty()
 
+# Initialization
+# if 'bot' not in st.session_state:
+#     st.sidebar.text("init...")
+#     st.session_state['bot'] = ChatGPT()
+#     init_prompt = "i'd like you to act like a snobby AI and tell me what you think of my structured data"
+#     init_answer = st.session_state['bot'].ask(init_prompt)
+
 # col?
 # Store the initial value of widgets in session state
 if "visibility" not in st.session_state:
@@ -191,24 +248,30 @@ col1, col2 = st.columns(2)
 
 with col1:
     st.subheader("Input")
-    # st.radio(
-    #     "Set selectbox label visibility 👉",
-    #     key="visibility",
-    #     options=["visible", "hidden", "collapsed"],
-    # )
-    # st.json({ })
+
+    sys_name = st.text_area(
+        label="Name of the system in question",
+        value=option,
+        height=1,
+    )
+
     st_text = st.text_area(
-        label="Type in some text",
-        value="SELECT shipping FROM users WHERE shipping = '201 Thayer St Providence RI 02912'"
-        "\n\n"
-        "{user: Willie Porter, ip: 192.168.2.80, email: willie@gmail.com}",
-        height=200,
+        label= "Structured text used as input",
+        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" } }""",
+        # value="SELECT shipping FROM users WHERE shipping = '201 Thayer St Providence RI 02912'"
+        # "\n\n"
+        # "{user: Willie Porter, ip: 192.168.2.80, email: willie@gmail.com}",
+        height=300,
     )
+    button = st.button("Detect and replace PII")
+    st.text("""""")
 
 with col2:
-    st.subheader("Analyzed")
+    st.subheader("Analyzed results with detected entities highlighted")
+    # st.text("Output text with detected entities highlighted")
     with st.spinner("Analyzing..."):
         if button or st.session_state.first_load:
+            option = sys_name
             st_analyze_results = analyze(
                 text=st_text,
                 entities=st_entities,
@@ -216,49 +279,212 @@ with col2:
                 score_threshold=st_threshold,
                 return_decision_process=st_return_decision_process,
             )
+            # """
+            # Ugly hack that checks if last 2 chars as Z" and changes the end of the last entity to -1
+            # This is done to prevent the anotation to inlcude the quotes for the date4 and breka the json donwtheroad
+            # ### TODO: make this less hacky?
+            # """
+            for i in st_analyze_results:
+                # st.write(i)
+                # st.write(st_text[i.end - 2:i.end])
+                if st_text[i.end-2:i.end] == 'Z"':# and i.type == "DATE_TIME":
+                    i.end = i.end-1
+                    continue
+                if st_text[i.end-2:i.end] == "Z'":# and i.type == "DATE_TIME":
+                    i.end = i.end-1
+                    continue
+                # if "'" in st_text[i.start:i.end]:
+                #     st_analyze_results.remove(i)
+                #     continue
+                # if "," in st_text[i.start:i.end]:
+                #     st_analyze_results.remove(i)
+                #     continue
+
+
             annotated_tokens = annotate(st_text, st_analyze_results, st_entities)
             # annotated_tokens
             annotated_text(*annotated_tokens)
 
-# end of col
+    # vertical space
+    st.text("")
+    st.text("")
 
-if 'first_load' not in st.session_state:
-    st.session_state['first_load'] = True
+    with st.expander("Show results with replaced PII and detailed results"):
+        # st.subheader("Final results with tokens instead if PII")
+        # vertical space
 
-# After
-
-st.subheader("Anonymized final results")
-
-with st.spinner("Anonymizing..."):
-    if button or st.session_state.first_load:
-        st_anonymize_results = anonymize(st_text, st_analyze_results)
-        st_anonymize_results
-
-
-# table result
-st.subheader("Detailed Findings")
-if st_analyze_results:
-    res_dicts = [r.to_dict() for r in st_analyze_results]
-    for d in res_dicts:
-        d['Value'] = st_text[d['start']:d['end']]
-    df = pd.DataFrame.from_records(res_dicts)
-    df = df[["entity_type", "Value", "score", "start", "end"]].rename(
-        {
-            "entity_type": "Entity type",
-            "start": "Start",
-            "end": "End",
-            "score": "Confidence",
-        },
-        axis=1,
-    )
+        if button or st.session_state.first_load:
+            st_anonymize_results = anonymize(st_text, st_analyze_results)
+            st.write(st_anonymize_results)
+            # st.write(st_anonymize_results)
+            # try:
+            #     # st_anonymize_results = ast.literal_eval(st_anonymize_results)
+            #     st.json(st_anonymize_results) #.replace("'", '"'))
+            # except Json Parse Error as e:
+            #     st.write(st_anonymize_results)
+        # vertical space
+        st.text("")
+        st.subheader("Detailed Findings")
+        if st_analyze_results:
+            res_dicts = [r.to_dict() for r in st_analyze_results]
+            for d in res_dicts:
+                d['Value'] = st_text[d['start']:d['end']]
+            df = pd.DataFrame.from_records(res_dicts)
+            df = df[["entity_type", "Value", "score", "start", "end"]].rename(
+                {
+                    "entity_type": "Entity type",
+                    "start": "Start",
+                    "end": "End",
+                    "score": "Confidence",
+                },
+                axis=1,
+            )
 
-    st.dataframe(df, width=1000)
-else:
-    st.text("No findings")
+            st.dataframe(df, width=1000)  # , height=500)
+        else:
+            st.text("No findings")
+
+# st_analyze_results
+# end of col
 
+# After the columns
+
+col5, col6 = st.columns(2)
+prompt = "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 " \
+            "strings and tokens in angular-brackets with an approximations to make it more human readable: \"{}\" ".format(
+        option, option,
+        st_anonymize_results)
+
+with col5:
+    st.subheader("Formatting")
+
+    button_create = st.button("Create summary")
+    st.markdown(
+        "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.")
+    st.text("""""")
+
+    with st.expander("Additional inputs"):
+        st_prompt = st.text_area(
+            label="Tokenized input with the formatted prompt",
+            value=prompt,
+            height=200,
+        )
+
+        write_results = ""
+        st_output = st.text_area(
+            label="Record results for later use",
+            value=write_results,
+            height=100,
+        )
+        button_save = st.button("Save summary to file?")
+        st.text("""""")
+
+        placeholder_table = st.empty()
+        placeholder_table.write("")
+
+        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."""
+        st_init_prompt = st.text_area(
+            label="Initial promopt to focus model",
+            value=init_prompt,
+            height=100,
+        )
+        button_reset = st.button("Reset model setup")
+
+import random
+
+with col6:
+    st.subheader("Output incident summary")
+
+# effect button_create(button2)
+# with st.spinner("button_create..."):
+if button_create:
+    # load existing promp and results
+    if (not api_togg):
+        saved_prompt = dataframe[dataframe['Title'] == option]['Prompt'].values[0]
+        saved_result = dataframe[dataframe['Title'] == option]['Result'].values[0]
+    else:
+        saved_prompt = ""
+        saved_result = ""
+    # check if match to current prompt
+    # if re.sub(r"[\n\t\s]*", "", saved_prompt) == re.sub(r"[\n\t\s]*", "", st_prompt):
+    md_results = ""
+    with col6:
+        x = st.empty()
+        x.markdown("")
+
+    # check if saved_prompt is not of a type float
+    if (not isinstance(saved_prompt, float)) and (not api_togg):
+        # st.write(saved_prompt)
+        with col6:
+            with st.spinner('Fetching results...'):
+                time.sleep(random.uniform(2.1, 5.8))
+                # st.write("Prompt already queried in the past, loading result from database")
+                md_results = saved_result
+                words = md_results.split()
+                num_words = len(words)
+                chunk_size = int(random.uniform(2, 6))
+                str_placeholder = ""
+
+                for i in range(0, num_words, chunk_size):
+                    chunk = ' '.join(words[i:i + chunk_size])
+                    str_placeholder = str_placeholder + " " + chunk
+                    x.markdown(str_placeholder)
+                    # x.markdown(chunk)
+                    time.sleep(random.uniform(0.1, 0.6))
+                x.markdown(saved_result)
+    else:
+        # st.write("New prompt, need GPT")
+        with col6:
+            with st.spinner('Generating, please wait...'):
+                bot = ChatGPT()
+                # init_answer = bot.ask(init_prompt)
+
+                init_points = ""
+                for chunk in bot.ask_stream(init_prompt):
+                    init_points = init_points + "."
+                    x.markdown(init_points)
+
+                x.markdown("")
+
+                # st_prompt = "tell me two facts about yourself"
+                for chunk in bot.ask_stream(st_prompt):
+                    md_results = md_results + chunk
+                    x.markdown(md_results)
+                    #check if last char of chunk is a new line
+                    # if "\n" in chunk:
+                    #     x.markdown(md_results)
+                    # st.markdown(chunk)
+                x.markdown(md_results)
+                bot._cleanup()
+                # md_results = bot.ask(st_prompt) #"Hello, could you tell what is {}?".format(option))
+                # print(md_results)  # prints the response from chatGPT
+
+                # st.write(st_prompt)
+                # st.write(saved_prompt)
+            # md_results = """No result found""" ##here GPT
+    # with col6:
+    #     # st.subheader("Output incident summary")
+    #     st.markdown(md_results)
+    placeholder_table.write((dataframe.loc[dataframe['Title'] == option]))
+
+# if button_reset:
+#     bot = ChatGPT()
+#     bot._cleanup()
+
+if button_save:
+    # dataframe = pd.read_csv("structured-data-anonymizer/assets/df_data_short.csv")
+    # save st_prompt and st_output to dataframe in row for Title = json_dict_option
+    dataframe.loc[dataframe['Title'] == option, 'Prompt'] = st_prompt
+    dataframe.loc[dataframe['Title'] == option, 'Result'] = md_results #st_output
+    # st.write(json_dict_option)
+    # write dataframe back to the csv file
+    dataframe.to_csv("structured-data-anonymizer/assets/df_data_short.csv", index=False)
+    st.write("Saved to file")
+    st.write(dataframe.loc[dataframe['Title'] == option])
+
+# end of document
 st.session_state['first_load'] = True
 
-# json result
 
 
 class ToDictListEncoder(JSONEncoder):