import pandas as pd
from openpyxl import load_workbook
from openpyxl.styles import PatternFill
from openpyxl.utils.dataframe import dataframe_to_rows

# Load both Excel files using pandas
base_file = 'base.xlsx'  # Excel file that contains the "Details" (CVE) and "Base Score" columns
target_file = 'target.xlsx'  # Excel file that contains "CVE Number" and hyperlinks

# Load the data into pandas dataframes (without affecting the hyperlinks)
df_base = pd.read_excel(base_file)  # This file contains 'Details' and 'Base Score' columns
df_target = pd.read_excel(target_file)  # This file contains the 'CVE Number' column and other data

# Load the target workbook with openpyxl to extract the hyperlinks
wb_target = load_workbook(target_file)
ws_target = wb_target.active

# Create a dictionary to store the hyperlinks
hyperlinks = {}
for row in ws_target.iter_rows(min_row=2, max_row=ws_target.max_row, min_col=1, max_col=ws_target.max_column):
    cve_cell = row[0]  # Assuming "CVE Number" is in the first column
    if cve_cell.hyperlink:
        hyperlinks[cve_cell.value] = cve_cell.hyperlink.target

# Merge based on 'Details' from df_base and 'CVE Number' from df_target
merged_df = pd.merge(df_target, df_base[['Details', 'Base Score']], left_on='CVE Number', right_on='Details', how='left')

# Write headers manually
header = list(merged_df.columns)
for col_num, col_name in enumerate(header, start=1):
    ws_target.cell(row=1, column=col_num, value=col_name)  # Set header row

# Write the merged data back into the Excel sheet, while restoring hyperlinks
for r_idx, row in enumerate(dataframe_to_rows(merged_df, index=False, header=False), start=2):
    for c_idx, value in enumerate(row, start=1):
        ws_target.cell(row=r_idx, column=c_idx, value=value)  # Set the cell value
    # Restore hyperlink for CVE Number if it exists
    cve_value = row[0]  # Assuming "CVE Number" is the first column in the merged data
    if cve_value in hyperlinks:
        ws_target.cell(row=r_idx, column=1).hyperlink = hyperlinks[cve_value]  # Restore the hyperlink

# Add conditional formatting
# Define fills for different conditions
yellow_fill = PatternFill(start_color="FFFF00", end_color="FFFF00", fill_type="solid")
orange_fill = PatternFill(start_color="FFA500", end_color="FFA500", fill_type="solid")
red_fill = PatternFill(start_color="FF0000", end_color="FF0000", fill_type="solid")

# Assuming 'Base Score' is in the third column of the merged data (adjust as
# necessary)
base_score_col = header.index('Base Score') + 1

# Apply conditional formatting to the 'Base Score' column
for row in ws_target.iter_rows(min_row=2, min_col=base_score_col, max_col=base_score_col, max_row=ws_target.max_row):
    for cell in row:
        try:
            # Convert cell value to a number
            cell_value = float(cell.value)
            if cell_value > 9:  # Apply format if 'Base Score' is greater than 9
                cell.fill = red_fill
        except (ValueError, TypeError):
            # Handle cells that cannot be converted to float (e.g., empty cells or non-numeric text)
            pass

# Add conditional formatting for 'CVE Title' containing "Remote Code Execution"
# Assuming 'CVE Title' is in the fourth column of the merged data (adjust as
# necessary)
cve_title_col = header.index('CVE Title') + 1

# Apply conditional formatting to the 'CVE Title' column
for row in ws_target.iter_rows(min_row=2, min_col=cve_title_col, max_col=cve_title_col, max_row=ws_target.max_row):
    for cell in row:
        if cell.value and "Remote Code Execution" in str(cell.value):  # Apply format if 'CVE Title' contains "Remote Code Execution"
            cell.fill = red_fill

# Add conditional formatting for 'Vector String'
# Assuming 'Vector String' is in the column of the merged data (adjust as
# necessary)
vector_string_col = header.index('Vector String') + 1

# Apply conditional formatting to the 'Vector String' column
for row in ws_target.iter_rows(min_row=2, min_col=vector_string_col, max_col=vector_string_col, max_row=ws_target.max_row):
    for cell in row:
        if cell.value:
            value_str = str(cell.value)
            if "AV:N" in value_str and "UI:N" in value_str and "PR:N" in value_str:
                cell.fill = red_fill
            elif "AV:N" in value_str and "UI:N" in value_str:
                cell.fill = orange_fill
            elif "AV:N" in value_str:
                cell.fill = yellow_fill

# Save the workbook with hyperlinks and conditional formatting intact
wb_target.save('merged_file_with_hyperlinks_and_formatting.xlsx')

print("Merge successful with conditional formatting. The new file is 'merged_file_with_hyperlinks_and_formatting.xlsx'.")

# Wait for input before closing
input("Press Enter to exit...")