#    Copyright (C) 2025 James Brotosky, Brandon Wickline
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.

#Standard Libary Imports:
import os
import pandas as pd

def colorText(text: str, color: str) -> str:
    colors = {
        "red": "\033[91m",
        "green": "\033[92m",
        "yellow": "\033[93m",
        "blue": "\033[94m",
        "magenta": "\033[95m",
        "cyan": "\033[96m",
        "white": "\033[97m",
        "reset": "\033[0m"
    }
    
    return f"{colors.get(color, colors['reset'])}{text}{colors['reset']}"

def style_dataframe_dark(df, output_html_path=None, overwrite=True):
    from datetime import datetime

    # Get current date and filename for subtitle
    today = datetime.now().strftime("%d %B %Y")  # Changed to "Day Month Year"
    filename = output_html_path.replace('.html', '') if output_html_path else "Report"

    dark_css = """
    <style>
    body {
        background-color: #000000;
        margin: 0;
        padding: 0;
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
        color: #f8f8f2;
    }
    .header {
        text-align: center;
        margin: 20px auto;
        padding: 10px;
        border-bottom: 2px solid #ffd700;
        max-width: 95%;
    }
    .header h1 {
        color: #ffd700;
        margin: 0;
        font-size: 32px;
    }
    .header p {
        color: #00bfff;
        margin: 5px 0 0 0;
        font-size: 18px;
    }
    .table-container {
        overflow-y: scroll;
        margin: 0 auto;
        width: 95%;
        max-height: calc(80vh - 100px);
        display: block;
        border: 1px solid #3a3a4d;
        margin-bottom: 0;
    }
    table {
        border-collapse: collapse;
        font-size: 14px;
        background-color: #1e1e2f;
        color: #f8f8f2;
        width: max-content;
    }
    th, td {
        border: 1px solid #3a3a4d;
        text-align: left;
        padding: 10px;
        max-width: 300px;
        word-wrap: break-word;
        overflow-wrap: break-word;
    }
    /* First column: no wrap */
    td:nth-child(1), th:nth-child(1) {
        white-space: nowrap;
        max-width: none !important;
        word-wrap: normal !important;
    }
    th {
        background-color: #2e2e40;
        color: #ffd700;
        position: sticky;
        top: 0;
        z-index: 10;
    }
    tr:nth-child(even) {
        background-color: #262638;
    }
    tr:hover {
        background-color: #33334d;
        color: #00bfff;
    }
    /* Custom scrollbar styling */
    .table-container::-webkit-scrollbar {
        width: 12px;
    }
    .table-container::-webkit-scrollbar-track {
        background: #1e1e2f;
    }
    .table-container::-webkit-scrollbar-thumb {
        background-color: #3a3a4d;
        border-radius: 6px;
    }
    </style>
    """

    header = f"""
    <div class="header">
        <h1>Airlock Tools</h1>
        <p>{filename} - {today}</p>
    </div>
    """

    html_table = df.to_html(index=False, escape=False)
    styled_html = (
        f"<html>\n"
        f"<head><title>Airlock Tools Report</title></head>\n"
        f"<body>\n"
        f"{dark_css}\n"
        f"{header}\n"
        f"<div class='table-container'>\n"
        f"  {html_table}\n"
        f"</div>\n"
        f"</body>\n"
        f"</html>"
    )
    if output_html_path:
        with open(output_html_path, "w", encoding="utf-8") as f:
            f.write(styled_html)
        print(f"✅ Styled table saved to '{output_html_path}'")
    elif overwrite:
        import tempfile
        temp_path = tempfile.mktemp(suffix=".html")
        with open(temp_path, "w", encoding="utf-8") as f:
            f.write(styled_html)
        print(f"✅ Styled table saved to temporary file: {temp_path}")
    else:
        return styled_html


def displayIntro():
    
    print(colorText(r"""
                                                           ███                  
                ████                                  ░████████                 
               █████████████                      ███████████████               
              █████████████████████          █████████████████████              
              ███████████████████           ██████████████████████▓             
             ███████████████████             ██████████████████████             
            █████████████████████            ███████████████████████            
            ████████████████████████████████████████████████████████            
            █████████      ██              ██              █████████            
            █████████      ██      ███      █              █████████            
            █████████      ██      ████     █████      █████████████            
            █████████      ██              ██████      █████████████            
             ████████      ██             ███████      ████████████░            
              ███████      ██      ██▓     ██████      ████████████             
               ██████      ██      ████     █████      ███████████              
                █████████████████████████████████████████████████               
                  ▒████████████████████      ██████████████████                 
                     ███████████████████      ███████████████▒                  
                         ███████████████      █████████████                     
                                ██████████    ███████████                       
                                             ████████                           
                                             ████                               
""", "yellow"))
    print(colorText(r"""
   _____  .__       .__                 __     ___________           .__          
  /  _  \ |__|______|  |   ____   ____ |  | __ \__    ___/___   ____ |  |   ______
 /  /_\  \|  \_  __ \  |  /  _ \_/ ___\|  |/ /   |    | /  _ \ /  _ \|  |  /  ___/
/    |    \  ||  | \/  |_(  <_> )  \___|    <    |    |(  <_> |  <_> )  |__\___ \ 
\____|__  /__||__|  |____/\____/ \___  >__|_ \   |____| \____/ \____/|____/____  >
        \/                           \/     \/                                 \/ 
""", "cyan"))
    print(colorText("=================================================================================", "cyan"))
    print(colorText("======================== Welcome to the Airlock API Tool ========================", "cyan"))
    print(colorText("=================================================================================", "cyan"))

def printEnforceChecklist(parq_base_dir,appr_base_dir, needappr_base_dir, pflight_base_dir, policy_list, allowlist_name, destination_name):

    print(colorText("\n --------------------------------------------------------------------", "cyan"))
    print(colorText(" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------", "cyan"))
    print(colorText(" --------------------------------------------------------------------", "cyan"))
    print(colorText("\nSequentually follow these steps to prepare a policy for enforcement:", "white"))
  
    print(colorText("\n1. Choose which originating policy or policies to move to enforcement", "cyan"))
    if not policy_list:
        print(colorText(f"  [✗]  No policies have been chosen","red"))
    else: 
        print(colorText(f"The following policies have been choosen:", "green"))
        for policy in policy_list:
            print(colorText(f"  [✓] {policy}","green"))


    print(colorText(f"2. Choose the destination policy and  allowlist", "cyan"))    
  
    if allowlist_name == " ": 
        print(colorText(f"  [✗]  No allowlists have been chosen","red"))
    elif allowlist_name != " " and allowlist_name != " "  and allowlist_name is not allowlist_name:
        print(colorText(f"  [✓]  {allowlist_name} has been selected as allowlist","green"))

    if destination_name == " ":
        print(colorText(f"  [✗]  No destination policy has been chosen","red"))
    else:
        print(colorText(f"  [✓]  destination policy is {destination_name}","green"))

    print(colorText("3. Pull and stage event history, combine the histories, add hash info, then categorize the hashes", "cyan"))
    if not policy_list:
        print(colorText(f"  [✗]  No policies have been chosen","red"))
    else:
        for policy in policy_list:
            if os.path.exists(f"{parq_base_dir}Exec_Hist_{policy}.parquet"): print(colorText(f"  [✓] Data for {policy} has been fetched","green"))
            else: print(colorText(f"  [✗] Data for {policy} has not been fetched","red"))

    print(colorText(f"4. Manually review the files:","cyan"))
    print(colorText("     'prepare_policy\\needs_approved\\good_hashes.csv' and 'prepare_policy\\needs_approved\\unknown_hashes.csv'\n", "cyan"))
    print(colorText("     Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.", "cyan"))
    print(colorText("     If metarules need to be created, please make note of them, and remove the row from the csv.", "cyan"))
    print(colorText("     When complete, save both csv files to the directory 'approved' and choose this option.","cyan"))
    print(colorText("     This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed", "cyan"))

    if os.path.exists(f"{appr_base_dir}good_hashes.csv") and os.path.exists(f"{appr_base_dir}unknown_hashes.csv"):
        print(colorText("  [✓]  Reviewed hashes have been loaded","green"))
    else:
        print(colorText("  [✗]  Reviewed hashes have not been loaded","red"))

    if os.path.exists(f"{parq_base_dir}all_hashes.parquet"):
        print(colorText("  [✓]  The combined approved hashes list has been generated","green"))
    else:
        print(colorText("  [✗]  The combined approved hashes list has not been generated","red"))

    if os.path.exists(f"{needappr_base_dir}primary_Paths.csv"):
        print(colorText("  [✓]  Path review list created","green"))
    else:
        print(colorText("  [✗]  Path review list has not been created","red"))    


    print(colorText(f"5. Manually review the files \n     'prepare_policy\\needs_approved\\primary_Paths.csv'\n     'prepare_policy\\needs_approved\\secondary_Paths.csv'", "cyan"))
    print(colorText("     Remove the rows containing path exclusions you do not approve of. The secondary list can be not added at all if nothing is useful" , "cyan"))
    print(colorText("     When complete, save the csv file to the directory 'approved'", "cyan"))
    print(colorText("     Do the same process with the list of publishers forthe same directories", "cyan"))
    print(colorText("     Preflight Lists will be generated", "cyan"))

    if  os.path.exists(f"{appr_base_dir}primary_Paths.csv"):
        print(colorText("  [✓]  Reviewed path list detected","green"))
    else:
        print(colorText("  [✗]  Path review list has not been detected","red"))    
    
    if    os.path.exists(f"{pflight_base_dir}final_path_exclusions.csv"):
        print(colorText("  [✓]  Preflight Path Exclusion List has been generated","green"))
    else:
        print(colorText("  [✗]  Preflight Path Exclusion List has not been generated","red"))    
        
    if os.path.exists(f"{pflight_base_dir}final_hash_approvals.csv"):
        print(colorText("  [✓]  Preflight hash approval list has been generated","green"))
    else:
        print(colorText("  [✗]  Preflight hash approval list has not been generated","red"))    




    print(colorText(f"6. Test ------------------------------------------------------", "cyan"))
    print(colorText(f"     Print rather than apply selected data.", "cyan"))    
  
    print(colorText(f"7. Liftoff ------------------------------------------------------", "cyan"))
    print(colorText(f"     Apply path exclusions according to allowed and approved paths", "cyan"))    
    print(colorText(f"     Apply signed or attested hashes to Parent Allow List", "cyan"))
    print(colorText(f"     Apply approved, but unsigned hashes to the Child Allow List", "cyan"))

    print(colorText("R. Remove/Reset Generated data - will prompt to allow keeping execution history", "cyan"))

    print(colorText("Q. Quit", "cyan"))

def areYouSure():
    print(colorText(f"🛑****************************************************************************************************************************************🛑","red"))
    print(colorText(f"⚠️=========================================================================================================================================⚠️","yellow"))
    print(colorText(f"🛑========================================================================================================================================🛑","red"))
    print(colorText(f"⚠️-------------This program will now begin to make changes to the Airlock Console. Do you understand and agree to proceed? ----------------⚠️", "yellow"))
    print(colorText(f"🛑========================================================================================================================================🛑","red"))
    print(colorText(f"⚠️=========================================================================================================================================⚠️","yellow"))
    print(colorText(f"🛑****************************************************************************************************************************************🛑","red"))

def locked():

    print(colorText(r"""
        ████████████████████████████████████████████████████████████████        
      ███                                                              ██       
      ██                             ██████                            ███      
      ██                          ████████████                         ███      
      ██                        ████        ███                        ███      
      ██                        ███          ███                       ███      
      ██                        ███          ███                       ███      
      ██                     ▒████████████████████                     ███      
      ██                     ██████████████████████                    ███      
      ██                     ██████████████████████                    ███      
      ██                     ██████████████████████                    ███      
      ██                     ██████████████████████                    ███      
      ██                     ██████████████████████                    ███      
      ██                                                               ███      
      ███                                                              ███      
      ████████████████████████████████████████████████████████████████████            
      ▒██████████████████████████████████████████████████████████████████▒  
                                     ▒████                                      
                                     ▒████                                      
                  ▓██████████████████████████████████████████                   
                 █████████████████████████████████████████████░  
""", "yellow"))
    
def printDeviceEnforceChecklist():

    print(colorText("\n --------------------------------------------------------------------", "cyan"))
    print(colorText(" ------------- 🛠️ 🔒 Prepare to Enforce Policy 🛠️ 🔒 ------------------", "cyan"))
    print(colorText(" --------------------------------------------------------------------", "cyan"))
    print(colorText("\nSequentually follow these steps to prepare a policy for enforcement:", "white"))
  
    print(colorText("\n1. Choose which originating policy or policies to move to enforcement", "cyan"))
    print(colorText("2. Pull and stage event history, combine the histories, add hash info, then categorize the hashes", "cyan"))
    print(colorText(f"3. Manually review the files:","cyan"))
    print(colorText("     'needs_approved\\good_{first_policy}_{second_policy}.csv' and 'needs_approved\\unknown_{first_policy}_{second_policy}.csv'", "cyan"))
    print(colorText("     Remove the rows containing hashes you do not approve of, and those you would not approve of without metarules.", "cyan"))
    print(colorText("     If metarules need to be created, please make note of them, and remove the row from the csv.", "cyan"))
    print(colorText("     When complete, save both csv files to the directory 'approved' and choose this option.","cyan"))
    print(colorText("     This will combine these approved hashes with the automatically approved hashes and generate a list of paths to be reviewed", "cyan"))
    print(colorText(f"4. Manually review the file 'needs_approved\\paths_needing_review.csv'", "cyan"))
    print(colorText("     Remove the rows containing path exclusions you do not approve of" , "cyan"))
    print(colorText("     When complete, save the csv file to the directory 'approved'", "cyan"))
    print(colorText("     Do the same process with the list of publishers forthe same directories", "cyan"))
    print(colorText("     Preflight Lists will be generated", "cyan"))

    print(colorText(f"5. Choose the destination policy and parent and child allow list", "cyan"))    
 
    print(colorText(f"6. Test ------------------------------------------------------", "cyan"))
    print(colorText(f"     Print rather than apply selected data.", "cyan"))    
  
    print(colorText(f"7. Liftoff ------------------------------------------------------", "cyan"))
    print(colorText(f"     Apply path exclusions according to allowed and approved paths", "cyan"))    
    print(colorText(f"     Apply signed or attested hashes to Parent Allow List", "cyan"))
    print(colorText(f"     Apply approved, but unsigned hashes to the Child Allow List", "cyan"))

    print(colorText("R. Remove/Reset Generated data - will prompt to allow keeping execution history", "cyan"))

    print(colorText("Q. Quit", "cyan"))


    
def apivalidation(): 
    match os.getenv('APIKEY'):
        case '':
            print(colorText("Please add your API Key to the .env file", "red"))


def tryToReadCSV(csv):
    try:
        if not os.path.exists(csv):
            print(colorText(f"Error: File '{csv}' does not exist.", "red"))
            return pd.DataFrame()  # Return empty DataFrame if file doesn't exist

        df = pd.read_csv(csv)
        if df.empty:
            print(colorText("Error: CSV file has headers but no data rows.", "red"))
        else:
            print(colorText(f"Data loaded successfully from {csv}", "green"))
    except pd.errors.EmptyDataError:
        print(colorText("Notice: CSV file is completely empty (no headers, no data), falling back to empty frame", "white"))
        df = pd.DataFrame()  # Create an empty DataFrame as fallback
    return df


def tryToReadParquet(parquet):
    try:
        df = pd.read_parquet(parquet)
        if df.empty:
            print(colorText("Error: Parquet file has headers but no data rows.", "red"))
        else:
            print(colorText(f"Data loaded successfully from {parquet}", "green"))
    except pd.errors.EmptyDataError:
            print(colorText("Notice : Parquet file is completely empty (no headers, no data), falling back to empty frame", "white"))
            df = pd.DataFrame()  # Create an empty DataFrame as fallback
    return df

def deduplicate_list(lst):
    seen = set()
    return [x for x in lst if not (x in seen or seen.add(x))]
