asghar_cve.py

import streamlit as st
import pandas as pd
import subprocess
import base64
from PIL import Image

# Title
st.title('Vulnerability and RVTools extractor')

def get_base64(bin_file):
    with open(bin_file, 'rb') as f:
        data = f.read()
    return base64.b64encode(data).decode()

def set_background(png_file):
    bin_str = get_base64(png_file)
    page_bg_img = '''
    <style>
    .stApp {
    background-image: url("data:image/png;base64,%s");
    background-size: cover;
    }
    </style>
    ''' % bin_str
    st.markdown(page_bg_img, unsafe_allow_html=True)

set_background('C:\\scripts\\CVE_Extractor\\background.PNG')

# Upload Roberto's Excel file
uploaded_file1 = st.file_uploader("Upload Roberto's file", type=["xlsx"], key='file1')

if uploaded_file1:
    # Read the first Excel file
    df1 = pd.read_excel(uploaded_file1)

    # Check if columns 'Vulnerability ID' and 'Asset' exist in the first Excel file
    if 'Vulnerability ID' in df1.columns and 'Asset' in df1.columns:
        # Extract unique contents from columns 'Vulnerability ID' and 'Asset'
        unique_C = df1['Asset'].drop_duplicates()
        unique_H = df1['Vulnerability ID'].drop_duplicates()

        # Create a new DataFrame for the unique contents
        new_df1 = pd.DataFrame({'Asset': df1['Asset'], 'Vulnerability ID': df1['Vulnerability ID']})
        #st.write("Unique contents from the first file:", new_df1)
    else:
        st.error("Columns 'Vulnerability ID' and 'Asset' do not exist in the first uploaded file.")
        new_df1 = pd.DataFrame()

    # Add a button to run the RVTools script
    if st.button('Run RVTools'):
        script_path = r"C:\scripts\CVE_Extractor\rvthemtools.bat"
        try:
            result = subprocess.run(["cmd", "/c", script_path], capture_output=True, text=True)
            st.write(f"Command executed: cmd /c {script_path}")
            if result.returncode == 0:
                st.success("RVTools script executed successfully.")
                st.text("Output:\n" + result.stdout)
                
                # Load the Excel file generated by RVTools
                rvtools_file = r"C:\scripts\CVE_Extractor\vminfo.xlsx"
                try:
                    df2 = pd.read_excel(rvtools_file)
                    st.write("RVTools data loaded successfully.")
                    #st.write("Columns in the RVTools file:", df2.columns)

                    # Check if required columns exist in the RVTools Excel file
                    required_columns = ['VM', 'Powerstate', 'OS according to the VMware Tools', 'Owner']
                    if all(col in df2.columns for col in required_columns):
                        # Drop rows with NaN values in the required columns
                        df2 = df2.dropna(subset=required_columns)
                        
                        # Delete all rows with VM names ending with '_DR'
                        df2_filtered = df2[~df2['VM'].str.endswith('_DR')]
                        #st.write("Filtered VMs (excluding _DR):", df2_filtered)

                        # Further filter rows where OS does not start with 'Microsoft'
                        df2_filtered = df2_filtered[df2_filtered['OS according to the VMware Tools'].str.startswith('Microsoft')]
                        #st.write("Filtered VMs (OS starts with 'Microsoft'):", df2_filtered)

                        # Filter and extract entries with powerstate 'powered on'
                        df2_powered_on = df2_filtered[df2_filtered['Powerstate'] == 'poweredOn']
                        #st.write("Filtered VMs (poweredOn):", df2_powered_on)

                        # Merge dataframes on Asset (from df1) and VM (from df2)
                        final_df = pd.merge(new_df1, df2_powered_on, left_on='Asset', right_on='VM')
                        final_df = final_df[['VM', 'OS according to the VMware Tools', 'Owner', 'Vulnerability ID']]
                        #st.write("Final result:", final_df)

                        # Add the new DataFrames to a new tab in the modified Excel file
                        modified_file = "modified_combined.xlsx"
                        with pd.ExcelWriter(modified_file, engine='openpyxl') as writer:
                            df1.to_excel(writer, sheet_name='Original File 1', index=False)
                            new_df1.to_excel(writer, sheet_name='Unique Contents File 1', index=False)
                            df2.to_excel(writer, sheet_name='Original File 2', index=False)
                            df2_powered_on.to_excel(writer, sheet_name='Filtered VMs File 2', index=False)
                            final_df.to_excel(writer, sheet_name='Matched Results', index=False)
                        
                        #st.success("Unique contents extracted and added to new tabs in the modified Excel file.")

                        # Allow the user to download the modified file
                        with open(modified_file, "rb") as file:
                            btn = st.download_button(
                                label="Download modified Excel file",
                                data=file,
                                file_name=modified_file,
                                mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                            )
                    else:
                        st.error("Required columns do not exist in the RVTools Excel file.")
                except Exception as e:
                    st.error(f"Error loading RVTools Excel file: {e}")
            else:
                st.error(f"Error executing RVTools script: {result.stderr}")
                st.text("Error Output:\n" + result.stderr)
        except Exception as e:
            st.error(f"Exception occurred: {e}")