From 3dfaec5618582e045e7dda2e146212aa64065318 Mon Sep 17 00:00:00 2001
From: kspviswa <7476271+kspviswa@users.noreply.github.com>
Date: Sun, 24 Mar 2024 20:12:38 -0500
Subject: [PATCH] exception handling and cosmetic changes

---
 bin/lpw_main.py   | 13 +++++++------
 bin/lpw_packet.py | 18 ++++++++++++------
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/bin/lpw_main.py b/bin/lpw_main.py
index 0319430..01ef181 100644
--- a/bin/lpw_main.py
+++ b/bin/lpw_main.py
@@ -6,6 +6,8 @@
 
 st.set_page_config(page_title='Local Packet Whisperer', page_icon='🗣️')
 
+lpw_avatar = "https://raw.githubusercontent.com/kspviswa/local-packet-whisperer/main/gifs/lpw_logo_small.png"
+
 DEFAULT_SYSTEM_MESSAGE = """
         You are a helper assistant specialized in analysing packet captures used to troubleshooting & technical analysis. Use the information present in packet_capture_info to answer all the questions truthfully. If the user asks about a specific application layer protocol, use the following hints to inspect the packet_capture_info to answer the question.
         
@@ -45,8 +47,7 @@ def save_sm() -> None:
 with st.sidebar:
     #st.markdown('## Enter OLLAMA endpoint:')
     #st.text_input("Ollama endpoint", placeholder="Eg localhost", label_visibility='hidden')
-    st.image(image="https://raw.githubusercontent.com/kspviswa/local-packet-whisperer/main/gifs/lpw_logo_small.png", width=400)
-    st.markdown('# Settings ⚙️')
+    st.markdown('# LPW Settings ⚙️', unsafe_allow_html=True)
     st.markdown('## Syster Message:')
     with st.expander(label='**Set System Message** *(optional)*', expanded=False):
         st.session_state['system_message'] = st.text_area(label='Override system message', value=st.session_state['system_message'], label_visibility="hidden", height=500, on_change=save_sm)
@@ -54,7 +55,7 @@ def save_sm() -> None:
     st.markdown('## Select a local model to use:')
     selected_model = st.selectbox('Models', placeholder="Choose an Option", options=options)
     if selected_model != "":
-        st.markdown(f"Selected :rainbow[{selected_model}]")
+        st.markdown(f"### :rainbow[Selected {selected_model}]")
     st.markdown('## Select protocols to filter in analysis')
     http = st.checkbox("HTTP") #80
     snmp = st.checkbox("SNMP") #161, 162
@@ -104,7 +105,7 @@ def getFiltersAndDecodeInfo():
 with col1:
     st.markdown('# :rainbow[Local Packet Whisperer] \n # :rainbow[🗣️🗣️🗣️]')
 with col2:
-    st.image(image="https://raw.githubusercontent.com/kspviswa/local-packet-whisperer/main/gifs/lpw_logo_small.png", use_column_width=True)
+    st.image(image=lpw_avatar, use_column_width=True)
 
 st.markdown('#### Step 1️⃣ 👉🏻 Build a knowledge base')
 packetFile = st.file_uploader(label='Upload either a PCAP or PCAPNG file to chat', accept_multiple_files=False, type=['pcap','pcapng'])
@@ -118,7 +119,7 @@ def getFiltersAndDecodeInfo():
         filters, decodes = getFiltersAndDecodeInfo()
         initLLM(pcap_data=getPcapData(input_file=f'{packetFile.name}', filter=filters, decode_info=decodes))
         os.remove(f'{packetFile.name}')
-    with st.chat_message(name='assistant'):
+    with st.chat_message(name='assistant', avatar=lpw_avatar):
         st.markdown('Chat with me..')
     for message in st.session_state.messages:
         with st.chat_message(name=message['role']):
@@ -127,7 +128,7 @@ def getFiltersAndDecodeInfo():
         st.session_state.messages.append({'role' : 'user', 'content' : prompt})
         with st.chat_message(name='user'):
             st.markdown(prompt)
-        with st.chat_message(name='assistant'):
+        with st.chat_message(name='assistant', avatar=lpw_avatar):
             with st.spinner('Processing....'):
                 full_response = chatWithModel(prompt=prompt, model=selected_model)
                 st.session_state.messages.append({'role' : 'assistant', 'content' : full_response})
diff --git a/bin/lpw_packet.py b/bin/lpw_packet.py
index 0c0c878..db88041 100644
--- a/bin/lpw_packet.py
+++ b/bin/lpw_packet.py
@@ -1,4 +1,5 @@
 import pyshark as ps
+import streamlit as st
 import os
 import re
 
@@ -12,10 +13,15 @@ def remove_ansi_escape_sequences(input_string):
     return cleaned_string
 
 def getPcapData(input_file:str = "", filter="", decode_info={}):
-    cap : ps.FileCapture = ps.FileCapture(input_file=input_file, display_filter=filter)
-    with open('out.txt', 'w') as f:
-        for pkt in cap:
-            print(pkt, file=f)
-    out_string = open('out.txt', 'r').read()
-    os.remove('out.txt')
+    try :
+        cap : ps.FileCapture = ps.FileCapture(input_file=input_file, display_filter=filter)
+        with open('out.txt', 'w') as f:
+            for pkt in cap:
+                print(pkt, file=f)
+        out_string = open('out.txt', 'r').read()
+        #os.remove('out.txt')
+    except ps.tshark.tshark.TSharkNotFoundException:
+        st.error(body='TShark/Wireshark is not installed. \n Please install [wireshark](https://tshark.dev/setup/install/#install-wireshark-with-a-package-manager) first', icon='🚨')
+        st.warning(body='LPW is now stopped', icon='🛑')
+        st.stop()
     return remove_ansi_escape_sequences(out_string)
\ No newline at end of file