@echo off
setlocal EnableExtensions EnableDelayedExpansion

REM =====================================================
REM   ____                  _       ____ ____  _   _ 
REM  |  _ \ __ _ _ ____   _( )___  / ___|  _ \| | | |
REM  | |_) / _` | '__\ \ / /|// __|| |   | | | | | | |
REM  |  __/ (_| | |   \ V /  \__ \| |___| |_| | |_| |
REM  |_|   \__,_|_|    \_/   |___(_)____|____/ \___/ 
REM                                                   
REM  Parv's Case Download Utility (PCDU) v1.3
REM  Downloads case files from analysis server
REM =====================================================

set "PCDU_VERSION=1.3"
set "PCDU_NAME=pcdu"
set "PCDU_INSTALL_DIR=%USERPROFILE%\PCDU"
set "SCRIPT_PATH=%~f0"
set "SCRIPT_NAME=%~nx0"

REM Icon URLs
set "ICON_URL_SJLNT=https://git.4parv.in/parv.ashwani/PCDU/raw/ASSETS/Icons/ServerIcon.ico"
set "ICON_URL_CASE=https://git.4parv.in/parv.ashwani/PCDU/raw/ASSETS/Icons/Casefolders.ico"
set "ICON_NAME_SJLNT=.SJLNT_Icon.ico"
set "ICON_NAME_CASE=.CaseFolder_Icon.ico"

REM =====================================================
REM Command Line Argument Handling
REM =====================================================
if /i "%~1"=="/install" goto :INSTALL_TO_PATH
if /i "%~1"=="/uninstall" goto :UNINSTALL_FROM_PATH
if /i "%~1"=="/help" goto :SHOW_HELP
if /i "%~1"=="-h" goto :SHOW_HELP
if /i "%~1"=="--help" goto :SHOW_HELP
if /i "%~1"=="/?" goto :SHOW_HELP
if /i "%~1"=="/update-icons" goto :UPDATE_ALL_ICONS

REM =====================================================
REM Check if first run (not installed to PATH)
REM =====================================================
call :CHECK_IF_INSTALLED
if !PCDU_INSTALLED! EQU 0 (
    call :PROMPT_INSTALL
)

goto :MAIN_SCRIPT

REM =====================================================
REM SET FOLDER ICON FUNCTION
REM Parameters: %1 = Folder Path, %2 = Icon URL, %3 = Icon Filename
REM =====================================================
:SET_FOLDER_ICON
setlocal EnableDelayedExpansion
set "FOLDER_PATH=%~1"
set "ICON_URL=%~2"
set "ICON_NAME=%~3"
set "ICON_PATH=%FOLDER_PATH%\%ICON_NAME%"
set "DESKTOP_INI=%FOLDER_PATH%\desktop.ini"

REM Check if folder exists
if not exist "%FOLDER_PATH%" (
    echo         [!!] Folder does not exist: %FOLDER_PATH%
    endlocal
    goto :EOF
)

REM Download icon if not present or force update
if not exist "%ICON_PATH%" (
    echo         Downloading icon...
    
    REM Try PowerShell first (more reliable)
    powershell -Command "try { [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12; Invoke-WebRequest -Uri '%ICON_URL%' -OutFile '%ICON_PATH%' -UseBasicParsing } catch { exit 1 }" >nul 2>&1
    
    if !ERRORLEVEL! NEQ 0 (
        REM Fallback to curl
        curl -sL -o "%ICON_PATH%" "%ICON_URL%" >nul 2>&1
    )
    
    if not exist "%ICON_PATH%" (
        echo         [!!] Failed to download icon
        endlocal
        goto :EOF
    )
    
    REM Hide the icon file
    attrib +h +s "%ICON_PATH%" >nul 2>&1
    echo         [OK] Icon downloaded
) else (
    echo         [OK] Icon already exists
)

REM Create or update desktop.ini
echo [.ShellClassInfo] > "%DESKTOP_INI%"
echo IconResource=%ICON_NAME%,0 >> "%DESKTOP_INI%"
echo IconFile=%ICON_NAME% >> "%DESKTOP_INI%"
echo IconIndex=0 >> "%DESKTOP_INI%"
echo InfoTip=PCDU Managed Folder >> "%DESKTOP_INI%"

REM Set desktop.ini attributes (hidden + system)
attrib +h +s "%DESKTOP_INI%" >nul 2>&1

REM Set folder as system folder (required for custom icon)
attrib +r "%FOLDER_PATH%" >nul 2>&1

echo         [OK] Folder icon configured

endlocal
goto :EOF

REM =====================================================
REM UPDATE ALL ICONS FUNCTION (for existing folders)
REM =====================================================
:UPDATE_ALL_ICONS
echo.
echo  ============================================================
echo   ^|^> PCDU Icon Update Utility
echo  ============================================================
echo.
echo   This will update icons for all existing PCDU folders.
echo.

REM Find all SJLNT folders
echo   Searching for SJLNT folders...
echo.

set "UPDATED_COUNT=0"

REM Check common locations
for %%D in (C D E F G H) do (
    if exist "%%D:\SJLNT" (
        echo   Found: %%D:\SJLNT
        echo     Updating SJLNT folder icon...
        call :SET_FOLDER_ICON "%%D:\SJLNT" "%ICON_URL_SJLNT%" "%ICON_NAME_SJLNT%"
        
        REM Update all case folders inside
        echo     Scanning for case folders...
        for /d %%C in ("%%D:\SJLNT\*") do (
            echo       Updating: %%~nxC
            call :SET_FOLDER_ICON "%%C" "%ICON_URL_CASE%" "%ICON_NAME_CASE%"
            set /a UPDATED_COUNT+=1
        )
    )
)

REM Also check C:\Temp\SJLNT
if exist "C:\Temp\SJLNT" (
    echo   Found: C:\Temp\SJLNT
    echo     Updating SJLNT folder icon...
    call :SET_FOLDER_ICON "C:\Temp\SJLNT" "%ICON_URL_SJLNT%" "%ICON_NAME_SJLNT%"
    
    for /d %%C in ("C:\Temp\SJLNT\*") do (
        echo       Updating: %%~nxC
        call :SET_FOLDER_ICON "%%C" "%ICON_URL_CASE%" "%ICON_NAME_CASE%"
        set /a UPDATED_COUNT+=1
    )
)

echo.
echo  ============================================================
echo   Icon update complete!
echo   Case folders updated: %UPDATED_COUNT%
echo.
echo   NOTE: You may need to refresh Explorer (F5) or restart
echo         Explorer for icons to appear.
echo  ============================================================
echo.

REM Refresh icon cache
echo   Refreshing icon cache...
ie4uinit.exe -show >nul 2>&1
ie4uinit.exe -ClearIconCache >nul 2>&1

pause
goto :EOF

REM =====================================================
REM INSTALL TO PATH FUNCTION
REM =====================================================
:INSTALL_TO_PATH
echo.
echo  ============================================================
echo   ^|^> PCDU Installation
echo  ============================================================
echo.

REM Check if already installed
call :CHECK_IF_INSTALLED
if !PCDU_INSTALLED! EQU 1 (
    echo   [OK] PCDU is already installed and accessible from PATH.
    echo       Location: !PCDU_FOUND_PATH!
    echo.
    choice /C YN /M "Reinstall/Update PCDU?"
    if !ERRORLEVEL! EQU 2 (
        echo.
        echo   Installation cancelled.
        goto :EOF
    )
)

REM Create install directory
echo   [1/5] Creating installation directory...
if not exist "%PCDU_INSTALL_DIR%" (
    mkdir "%PCDU_INSTALL_DIR%" 2>nul
    if !ERRORLEVEL! NEQ 0 (
        echo   [ERROR] Failed to create directory: %PCDU_INSTALL_DIR%
        goto :EOF
    )
    echo         Created: %PCDU_INSTALL_DIR%
) else (
    echo         Exists: %PCDU_INSTALL_DIR%
)

REM Copy script to install directory
echo.
echo   [2/5] Copying PCDU script...
copy /Y "%SCRIPT_PATH%" "%PCDU_INSTALL_DIR%\%PCDU_NAME%.bat" >nul 2>&1
if !ERRORLEVEL! NEQ 0 (
    echo   [ERROR] Failed to copy script to installation directory.
    goto :EOF
)
echo         Copied to: %PCDU_INSTALL_DIR%\%PCDU_NAME%.bat

REM Create a simple launcher without .bat extension (optional convenience)
echo   [2b/5] Creating launcher...
echo @echo off > "%PCDU_INSTALL_DIR%\%PCDU_NAME%.cmd"
echo call "%PCDU_INSTALL_DIR%\%PCDU_NAME%.bat" %%* >> "%PCDU_INSTALL_DIR%\%PCDU_NAME%.cmd"
echo         Created: %PCDU_INSTALL_DIR%\%PCDU_NAME%.cmd

REM Download icons to install directory for future use
echo.
echo   [3/5] Downloading icons...
if not exist "%PCDU_INSTALL_DIR%\Icons" mkdir "%PCDU_INSTALL_DIR%\Icons"

powershell -Command "try { [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12; Invoke-WebRequest -Uri '%ICON_URL_SJLNT%' -OutFile '%PCDU_INSTALL_DIR%\Icons\ServerIcon.ico' -UseBasicParsing } catch { exit 1 }" >nul 2>&1
if exist "%PCDU_INSTALL_DIR%\Icons\ServerIcon.ico" (
    echo         [OK] ServerIcon.ico downloaded
) else (
    echo         [!!] Failed to download ServerIcon.ico
)

powershell -Command "try { [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12; Invoke-WebRequest -Uri '%ICON_URL_CASE%' -OutFile '%PCDU_INSTALL_DIR%\Icons\Casefolders.ico' -UseBasicParsing } catch { exit 1 }" >nul 2>&1
if exist "%PCDU_INSTALL_DIR%\Icons\Casefolders.ico" (
    echo         [OK] Casefolders.ico downloaded
) else (
    echo         [!!] Failed to download Casefolders.ico
)

REM Check if directory is already in PATH
echo.
echo   [4/5] Checking PATH environment variable...
set "PATH_CONTAINS_PCDU=0"
echo !PATH! | findstr /I /C:"%PCDU_INSTALL_DIR%" >nul 2>&1
if !ERRORLEVEL! EQU 0 (
    set "PATH_CONTAINS_PCDU=1"
    echo         PCDU directory already in PATH.
)

REM Add to PATH if not present
if !PATH_CONTAINS_PCDU! EQU 0 (
    echo         Adding PCDU to User PATH...
    
    REM Get current user PATH
    for /f "tokens=2*" %%A in ('reg query "HKCU\Environment" /v PATH 2^>nul') do set "USER_PATH=%%B"
    
    REM Check if USER_PATH is empty or not set
    if "!USER_PATH!"=="" (
        set "NEW_PATH=%PCDU_INSTALL_DIR%"
    ) else (
        set "NEW_PATH=!USER_PATH!;%PCDU_INSTALL_DIR%"
    )
    
    REM Use setx to permanently add to user PATH
    setx PATH "!NEW_PATH!" >nul 2>&1
    if !ERRORLEVEL! EQU 0 (
        echo         [OK] Added to User PATH successfully.
    ) else (
        echo         [WARNING] Could not add to PATH automatically.
        echo         Please manually add this to your PATH:
        echo         %PCDU_INSTALL_DIR%
    )
) 

REM Verify installation
echo.
echo   [5/5] Verifying installation...

REM Update current session PATH
set "PATH=%PATH%;%PCDU_INSTALL_DIR%"

if exist "%PCDU_INSTALL_DIR%\%PCDU_NAME%.bat" (
    echo         [OK] Script installed successfully.
) else (
    echo         [ERROR] Script not found at expected location.
    goto :EOF
)

echo.
echo  ============================================================
echo   ^|^> Installation Complete!
echo  ============================================================
echo.
echo   PCDU has been installed to: %PCDU_INSTALL_DIR%
echo.
echo   You can now run PCDU from anywhere using:
echo.
echo     pcdu                      - Run interactively
echo     pcdu [CaseNumber]         - Download specific case
echo     pcdu [CaseNo] [User]      - With username
echo     pcdu [CaseNo] [User] [Pw] - Fully automated
echo.
echo     pcdu /help                - Show help
echo     pcdu /install             - Reinstall PCDU
echo     pcdu /uninstall           - Remove PCDU from system
echo     pcdu /update-icons        - Update all folder icons
echo.
echo   NOTE: You may need to restart your terminal/CMD for
echo         PATH changes to take effect.
echo.
echo  ============================================================
echo.

REM Refresh environment for current process
call :REFRESH_ENV

pause
goto :EOF

REM =====================================================
REM UNINSTALL FROM PATH FUNCTION
REM =====================================================
:UNINSTALL_FROM_PATH
echo.
echo  ============================================================
echo   ^|^> PCDU Uninstallation
echo  ============================================================
echo.

choice /C YN /M "Are you sure you want to uninstall PCDU?"
if !ERRORLEVEL! EQU 2 (
    echo.
    echo   Uninstallation cancelled.
    goto :EOF
)

echo.
echo   [1/3] Removing PCDU from PATH...

REM Get current user PATH
for /f "tokens=2*" %%A in ('reg query "HKCU\Environment" /v PATH 2^>nul') do set "USER_PATH=%%B"

REM Remove PCDU_INSTALL_DIR from PATH
set "NEW_PATH=!USER_PATH:%PCDU_INSTALL_DIR%;=!"
set "NEW_PATH=!NEW_PATH:;%PCDU_INSTALL_DIR%=!"
set "NEW_PATH=!NEW_PATH:%PCDU_INSTALL_DIR%=!"

REM Update PATH
if not "!NEW_PATH!"=="!USER_PATH!" (
    setx PATH "!NEW_PATH!" >nul 2>&1
    echo         [OK] Removed from PATH.
) else (
    echo         [--] PCDU was not in PATH.
)

echo.
echo   [2/3] Removing installation directory...
if exist "%PCDU_INSTALL_DIR%" (
    rmdir /S /Q "%PCDU_INSTALL_DIR%" 2>nul
    if !ERRORLEVEL! EQU 0 (
        echo         [OK] Removed: %PCDU_INSTALL_DIR%
    ) else (
        echo         [WARNING] Could not remove directory. Please delete manually:
        echo         %PCDU_INSTALL_DIR%
    )
) else (
    echo         [--] Installation directory not found.
)

echo.
echo   [3/3] Cleanup complete.

echo.
echo   NOTE: Folder icons on SJLNT and case folders will remain.
echo         To remove them, delete desktop.ini and .ico files
echo         from those folders.
echo.
echo  ============================================================
echo   ^|^> PCDU has been uninstalled.
echo  ============================================================
echo.
echo   NOTE: Restart your terminal for PATH changes to take effect.
echo.

pause
goto :EOF

REM =====================================================
REM SHOW HELP FUNCTION
REM =====================================================
:SHOW_HELP
echo.
echo  ============================================================
echo   ^|^> Parv's Case Download Utility ^(PCDU^) v%PCDU_VERSION%
echo  ============================================================
echo.
echo   DESCRIPTION:
echo     Downloads case files from the analysis server including
echo     configurations, packet captures, and media files.
echo.
echo   USAGE:
echo     %PCDU_NAME% [options]
echo     %PCDU_NAME% [CaseNumber] [Username] [Password]
echo.
echo   OPTIONS:
echo     /install        Install PCDU to system PATH
echo     /uninstall      Remove PCDU from system
echo     /update-icons   Update folder icons for all SJLNT folders
echo     /help, -h       Show this help message
echo.
echo   EXAMPLES:
echo     %PCDU_NAME%                     Interactive mode
echo     %PCDU_NAME% 12345678            Download case 12345678
echo     %PCDU_NAME% 12345678 myuser     With username
echo     %PCDU_NAME% /install            Install to PATH
echo     %PCDU_NAME% /update-icons       Refresh all folder icons
echo.
echo   STORAGE LOCATIONS (in priority order):
echo     1. D:\SJLNT\[CaseNumber]
echo     2. [Temporary Storage Drive]:\SJLNT\[CaseNumber]
echo     3. C:\Temp\SJLNT\[CaseNumber] (fallback)
echo.
echo   FOLDER ICONS:
echo     - SJLNT folders get a Server icon
echo     - Case folders get a Case folder icon
echo     - Icons are downloaded from git.4parv.in
echo.
echo   DOWNLOADED FILES:
echo     - Saved configs (ns.conf) from collector bundles
echo     - Running configs (ns_running_config.conf)
echo     - Packet captures (.pcap, .cap, .har, etc.)
echo     - SSL keys (.sslkeys, .keys)
echo     - Media files (images)
echo     - Log and text files
echo.
echo   REQUIREMENTS:
echo     - PuTTY (pscp.exe, plink.exe) OR OpenSSH (scp.exe)
echo     - Network access to sjanalysis.citrite.net
echo.
echo  ============================================================
echo.
pause
goto :EOF

REM =====================================================
REM CHECK IF INSTALLED FUNCTION
REM =====================================================
:CHECK_IF_INSTALLED
set "PCDU_INSTALLED=0"
set "PCDU_FOUND_PATH="

REM Check if pcdu command is available
where pcdu.bat >nul 2>&1
if !ERRORLEVEL! EQU 0 (
    for /f "delims=" %%P in ('where pcdu.bat 2^>nul') do (
        set "PCDU_FOUND_PATH=%%P"
    )
    set "PCDU_INSTALLED=1"
    goto :EOF
)

where pcdu.cmd >nul 2>&1
if !ERRORLEVEL! EQU 0 (
    for /f "delims=" %%P in ('where pcdu.cmd 2^>nul') do (
        set "PCDU_FOUND_PATH=%%P"
    )
    set "PCDU_INSTALLED=1"
    goto :EOF
)

REM Check if running from install directory
if /i "%~dp0"=="%PCDU_INSTALL_DIR%\" (
    set "PCDU_INSTALLED=1"
    set "PCDU_FOUND_PATH=%PCDU_INSTALL_DIR%\%PCDU_NAME%.bat"
)

goto :EOF

REM =====================================================
REM PROMPT INSTALL FUNCTION
REM =====================================================
:PROMPT_INSTALL
echo.
echo  ============================================================
echo   ^|^> PCDU First Run Setup
echo  ============================================================
echo.
echo   PCDU is not installed to your system PATH.
echo   Installing allows you to run 'pcdu' from any directory.
echo.
echo   Installation will:
echo     - Create folder: %PCDU_INSTALL_DIR%
echo     - Copy PCDU script to that folder
echo     - Download folder icons
echo     - Add folder to your User PATH
echo.

choice /C YNC /M "Install PCDU to PATH? (Y=Yes, N=No, C=Cancel script)"
if !ERRORLEVEL! EQU 3 (
    echo.
    echo   Script cancelled.
    exit /b 0
)
if !ERRORLEVEL! EQU 1 (
    call :INSTALL_TO_PATH
    echo.
    echo   Continuing with case download...
    echo.
)

goto :EOF

REM =====================================================
REM REFRESH ENVIRONMENT FUNCTION
REM =====================================================
:REFRESH_ENV
REM This refreshes the PATH for the current cmd session
for /f "tokens=2*" %%A in ('reg query "HKCU\Environment" /v PATH 2^>nul') do set "PATH=%%B;%PATH%"
goto :EOF

REM =====================================================
REM SETUP SJLNT FOLDER ICON
REM =====================================================
:SETUP_SJLNT_ICON
echo.
echo  [ICON SETUP] Configuring SJLNT folder icon...
call :SET_FOLDER_ICON "%STORAGE_BASE%" "%ICON_URL_SJLNT%" "%ICON_NAME_SJLNT%"
goto :EOF

REM =====================================================
REM SETUP CASE FOLDER ICON
REM =====================================================
:SETUP_CASE_ICON
echo.
echo  [ICON SETUP] Configuring case folder icon...
call :SET_FOLDER_ICON "%DEST%" "%ICON_URL_CASE%" "%ICON_NAME_CASE%"
goto :EOF

REM =====================================================
REM MAIN SCRIPT
REM =====================================================
:MAIN_SCRIPT

REM =====================================================
REM Storage Drive Configuration
REM Priority: D: -> "Temporary Storage" -> C:\Temp
REM =====================================================
set "STORAGE_DRIVE="
set "STORAGE_BASE="
set "DRIVE_STATUS="

echo.
echo  [STORAGE CHECK] Checking available storage drives...
echo.

REM Check for D: drive first
if exist "D:\" (
    set "STORAGE_DRIVE=D:"
    set "STORAGE_BASE=D:\SJLNT"
    set "DRIVE_STATUS=D: Drive found and accessible"
    echo   [OK] D: Drive found and accessible
    goto :DRIVE_FOUND
)

echo   [--] D: Drive not found, checking for 'Temporary Storage' drive...

REM Check for drive with label "Temporary Storage"
for /f "tokens=1,2*" %%A in ('wmic logicaldisk get caption^,volumename 2^>nul ^| findstr /I "Temporary Storage"') do (
    set "TEMP_DRIVE=%%A"
    if exist "!TEMP_DRIVE!\" (
        set "STORAGE_DRIVE=!TEMP_DRIVE!"
        set "STORAGE_BASE=!TEMP_DRIVE!\SJLNT"
        set "DRIVE_STATUS=Temporary Storage drive found at !TEMP_DRIVE!"
        echo   [OK] 'Temporary Storage' drive found at !TEMP_DRIVE!
        goto :DRIVE_FOUND
    )
)

REM Alternative method: Check all drives for "Temporary Storage" label
for %%D in (E F G H I J K L M N O P Q R S T U V W X Y Z) do (
    if exist "%%D:\" (
        for /f "tokens=*" %%L in ('vol %%D: 2^>nul ^| findstr /I "Temporary Storage"') do (
            set "STORAGE_DRIVE=%%D:"
            set "STORAGE_BASE=%%D:\SJLNT"
            set "DRIVE_STATUS=Temporary Storage drive found at %%D:"
            echo   [OK] 'Temporary Storage' drive found at %%D:
            goto :DRIVE_FOUND
        )
    )
)

echo   [--] 'Temporary Storage' drive not found
echo.

REM Fallback to C:\Temp if no preferred drive found
echo   [!!] WARNING: Neither D: nor 'Temporary Storage' drive found!
echo.
echo   Available options:
echo     1. Use C:\Temp\SJLNT (fallback location)
echo     2. Cancel and check your drives
echo.
choice /C 12 /M "Select option"
if !ERRORLEVEL! EQU 2 (
    echo.
    echo   Operation cancelled by user.
    echo   Please ensure D: drive or a drive labeled 'Temporary Storage' is available.
    pause
    exit /b 1
)

set "STORAGE_DRIVE=C:"
set "STORAGE_BASE=C:\Temp\SJLNT"
set "DRIVE_STATUS=Using fallback location C:\Temp\SJLNT"
echo.
echo   [!!] Using fallback: C:\Temp\SJLNT

:DRIVE_FOUND

REM Check for pscp.exe (PuTTY SCP) or scp.exe
set "SCP_TOOL="
where pscp.exe >nul 2>&1
if !ERRORLEVEL! EQU 0 (
    set "SCP_TOOL=pscp.exe"
    set "SCP_TYPE=PSCP"
) else (
    where scp.exe >nul 2>&1
    if !ERRORLEVEL! EQU 0 (
        set "SCP_TOOL=scp.exe"
        set "SCP_TYPE=SCP"
    )
)

if "%SCP_TOOL%"=="" (
    echo.
    echo  ============================================================
    echo   [ERROR] SCP Tool Not Found
    echo  ============================================================
    echo.
    echo   Neither pscp.exe nor scp.exe found in PATH.
    echo.
    echo   Please install one of the following:
    echo     - PuTTY ^(includes pscp.exe^) - https://www.putty.org/
    echo     - OpenSSH Client ^(includes scp.exe^) - Built into Windows 10/11
    echo.
    echo   To enable OpenSSH on Windows:
    echo     Settings ^> Apps ^> Optional Features ^> Add OpenSSH Client
    echo.
    pause
    exit /b 1
)

REM =====================================================
REM Configuration
REM =====================================================
set "REMOTE_SERVER=sjanalysis.citrite.net"
set "REMOTE_BASE_PATH=/upload/ftp"

REM Prompt for case number if not provided
if "%~1"=="" (
    set /p CASENO=Enter Case Number: 
) else (
    REM Skip if argument is an option
    echo %~1 | findstr /B /C:"/" /C:"-" >nul
    if !ERRORLEVEL! EQU 0 (
        set /p CASENO=Enter Case Number: 
    ) else (
        set "CASENO=%~1"
    )
)

if "%CASENO%"=="" (
    echo ERROR: Case number cannot be empty.
    exit /b 1
)

REM Prompt for credentials if not provided
if "%~2"=="" (
    set /p USERNAME=Enter SSH Username: 
) else (
    set "USERNAME=%~2"
)

if "%USERNAME%"=="" (
    echo ERROR: Username cannot be empty.
    exit /b 1
)

if "%~3"=="" (
    echo Enter SSH Password for %USERNAME%@%REMOTE_SERVER%:
    REM Use PowerShell for hidden password input
    for /f "delims=" %%P in ('powershell -Command "$p = Read-Host -AsSecureString; [Runtime.InteropServices.Marshal]::PtrToStringAuto([Runtime.InteropServices.Marshal]::SecureStringToBSTR($p))"') do set "PASSWORD=%%P"
) else (
    set "PASSWORD=%~3"
)

if "%PASSWORD%"=="" (
    echo ERROR: Password cannot be empty.
    exit /b 1
)

REM =====================================================
REM Paths - Now using configured storage drive
REM =====================================================
set "REMOTE_CASE_PATH=%REMOTE_BASE_PATH%/%CASENO%"
set "DEST=%STORAGE_BASE%\%CASENO%"
set "LOGFILE=%DEST%\download_%CASENO%.log"
set "TEMP_LIST=%TEMP%\caselist_%CASENO%.txt"
set "TEMP_FILTERED=%TEMP%\filtered_%CASENO%.txt"
set "COLLECTOR_LIST=%TEMP%\collectors_%CASENO%.txt"
set "ALL_FILES_LIST=%TEMP%\allfiles_%CASENO%.txt"

echo.
echo  ============================================================
echo   ^|^> Parv's Case Download Utility ^(PCDU^) v%PCDU_VERSION%
echo  ============================================================
echo   Case Number    : %CASENO%
echo   Remote Server  : %REMOTE_SERVER%
echo   Remote Path    : %REMOTE_CASE_PATH%
echo   Storage Drive  : %STORAGE_DRIVE%
echo   Local Dest     : %DEST%
echo   SCP Tool       : %SCP_TYPE%
echo   User           : %USERNAME%
echo   Drive Status   : %DRIVE_STATUS%
echo  ============================================================
echo.

REM =====================================================
REM Create destination directory structure
REM =====================================================
echo  [PATH CHECK] Verifying destination paths...
echo.

REM Track if SJLNT folder was just created (for icon setup)
set "SJLNT_CREATED=0"
set "CASE_CREATED=0"

REM Check if base storage path exists
if exist "%STORAGE_BASE%" (
    echo   [OK] Base path exists: %STORAGE_BASE%
) else (
    echo   [!!] Base path does not exist: %STORAGE_BASE%
    echo       Creating base directory...
    mkdir "%STORAGE_BASE%" 2>nul
    if !ERRORLEVEL! EQU 0 (
        echo   [OK] Created: %STORAGE_BASE%
        set "SJLNT_CREATED=1"
    ) else (
        echo   [ERROR] Failed to create base directory: %STORAGE_BASE%
        echo.
        echo   Possible reasons:
        echo     - Drive not writable
        echo     - Insufficient permissions
        echo     - Drive full
        echo.
        pause
        exit /b 1
    )
)

REM Setup SJLNT folder icon (always ensure it's set)
call :SETUP_SJLNT_ICON

REM Check if case destination exists
if exist "%DEST%" (
    echo   [OK] Case path exists: %DEST%
    echo.
    echo   [!!] WARNING: Case folder already exists!
    echo       Files may be overwritten.
    echo.
    choice /C YN /M "Continue and potentially overwrite existing files?"
    if !ERRORLEVEL! EQU 2 (
        echo.
        echo   Operation cancelled by user.
        pause
        exit /b 0
    )
) else (
    echo   [!!] Case path does not exist: %DEST%
    echo       Creating case directory...
    mkdir "%DEST%" 2>nul
    if !ERRORLEVEL! EQU 0 (
        echo   [OK] Created: %DEST%
        set "CASE_CREATED=1"
    ) else (
        echo   [ERROR] Failed to create case directory: %DEST%
        pause
        exit /b 1
    )
)

REM Setup case folder icon
call :SETUP_CASE_ICON

REM Create subdirectories for organization
echo.
echo  [PATH CHECK] Creating subdirectories...

for %%S in (configs captures media) do (
    if exist "%DEST%\%%S" (
        echo   [OK] Exists: %DEST%\%%S
    ) else (
        mkdir "%DEST%\%%S" 2>nul
        if !ERRORLEVEL! EQU 0 (
            echo   [OK] Created: %DEST%\%%S
        ) else (
            echo   [!!] Failed to create: %DEST%\%%S
        )
    )
)

echo.

REM Initialize log file
echo ============================================================ > "%LOGFILE%"
echo  Parv's Case Download Utility ^(PCDU^) v%PCDU_VERSION% - Download Log >> "%LOGFILE%"
echo ============================================================ >> "%LOGFILE%"
echo  Case Number   : %CASENO% >> "%LOGFILE%"
echo  Date/Time     : %DATE% %TIME% >> "%LOGFILE%"
echo  Remote Server : %REMOTE_SERVER% >> "%LOGFILE%"
echo  User          : %USERNAME% >> "%LOGFILE%"
echo  Storage Drive : %STORAGE_DRIVE% >> "%LOGFILE%"
echo  Destination   : %DEST% >> "%LOGFILE%"
echo  Drive Status  : %DRIVE_STATUS% >> "%LOGFILE%"
echo ============================================================ >> "%LOGFILE%"
echo. >> "%LOGFILE%"

REM =====================================================
REM Test SSH connection and verify case directory
REM =====================================================
echo [1/5] Testing SSH connection...
echo [1/5] Testing SSH connection... >> "%LOGFILE%"

if "%SCP_TYPE%"=="PSCP" (
    echo y | plink.exe -batch -pw "%PASSWORD%" %USERNAME%@%REMOTE_SERVER% "test -d %REMOTE_CASE_PATH% && echo OK || echo NOTFOUND" 2>nul | findstr /V /C:"Keyboard" /C:"prompts from server" > "%TEMP_LIST%"
) else (
    sshpass -p "%PASSWORD%" ssh -o StrictHostKeyChecking=no -o UserKnownHostsFile=NUL %USERNAME%@%REMOTE_SERVER% "test -d %REMOTE_CASE_PATH% && echo OK || echo NOTFOUND" > "%TEMP_LIST%" 2>&1
)

if !ERRORLEVEL! NEQ 0 (
    echo ERROR: SSH connection failed. >> "%LOGFILE%"
    type "%TEMP_LIST%" >> "%LOGFILE%"
    echo ERROR: SSH connection failed.
    echo Check credentials and network connectivity.
    echo See log: %LOGFILE%
    del "%TEMP_LIST%" 2>nul
    pause
    exit /b 1
)

findstr /C:"OK" "%TEMP_LIST%" >nul
if !ERRORLEVEL! NEQ 0 (
    echo ERROR: Case directory not found on server: %REMOTE_CASE_PATH% >> "%LOGFILE%"
    echo ERROR: Case directory not found on server.
    echo Path: %REMOTE_CASE_PATH%
    echo Please verify the case number.
    del "%TEMP_LIST%" 2>nul
    pause
    exit /b 1
)

echo       Connection successful.
echo       Connection successful. >> "%LOGFILE%"

REM =====================================================
REM Single SSH call to get all required info
REM Collectors + All files in case root
REM =====================================================
echo.
echo [2/5] Scanning case directory...
echo [2/5] Scanning case directory... >> "%LOGFILE%"

if "%SCP_TYPE%"=="PSCP" (
    plink.exe -batch -pw "%PASSWORD%" %USERNAME%@%REMOTE_SERVER% "echo '===COLLECTORS==='; find %REMOTE_CASE_PATH% -maxdepth 1 -type d -name 'collector_*' 2>/dev/null; echo '===FILES==='; find %REMOTE_CASE_PATH% -maxdepth 1 -type f \( -name '*.cap' -o -name '*.pcap' -o -name '*.pcapng' -o -name '*.sslkeys' -o -name '*.keys' -o -name '*.har' -o -name '*.xml' -o -name '*.jpg' -o -name '*.jpeg' -o -name '*.png' -o -name '*.gif' -o -name '*.bmp' -o -name '*.webp' -o -name '*.tiff' -o -name '*.svg' -o -name '*.txt' -o -name '*.log' -o -name '*.csv' -o -name '*.json' \) 2>/dev/null" 2>nul | findstr /V /C:"Keyboard" /C:"prompts from server" > "%ALL_FILES_LIST%"
) else (
    sshpass -p "%PASSWORD%" ssh -o StrictHostKeyChecking=no %USERNAME%@%REMOTE_SERVER% "echo '===COLLECTORS==='; find %REMOTE_CASE_PATH% -maxdepth 1 -type d -name 'collector_*' 2>/dev/null; echo '===FILES==='; find %REMOTE_CASE_PATH% -maxdepth 1 -type f \( -name '*.cap' -o -name '*.pcap' -o -name '*.pcapng' -o -name '*.sslkeys' -o -name '*.keys' -o -name '*.har' -o -name '*.xml' -o -name '*.jpg' -o -name '*.jpeg' -o -name '*.png' -o -name '*.gif' -o -name '*.bmp' -o -name '*.webp' -o -name '*.tiff' -o -name '*.svg' -o -name '*.txt' -o -name '*.log' -o -name '*.csv' -o -name '*.json' \) 2>/dev/null" > "%ALL_FILES_LIST%"
)

REM Parse collectors from combined output
set "IN_COLLECTORS=0"
set "IN_FILES=0"
echo. > "%COLLECTOR_LIST%"
echo. > "%TEMP_LIST%"

for /f "usebackq delims=" %%A in ("%ALL_FILES_LIST%") do (
    set "LINE=%%A"
    
    REM Check for section markers
    echo !LINE! | findstr /C:"===COLLECTORS===" >nul
    if !ERRORLEVEL! EQU 0 (
        set "IN_COLLECTORS=1"
        set "IN_FILES=0"
    ) else (
        echo !LINE! | findstr /C:"===FILES===" >nul
        if !ERRORLEVEL! EQU 0 (
            set "IN_COLLECTORS=0"
            set "IN_FILES=1"
        ) else (
            REM Skip empty lines and filter noise
            if not "!LINE!"=="" (
                echo !LINE! | findstr /C:"Keyboard" /C:"prompts" /C:"authentication" >nul
                if !ERRORLEVEL! NEQ 0 (
                    if !IN_COLLECTORS! EQU 1 (
                        echo !LINE!>> "%COLLECTOR_LIST%"
                    )
                    if !IN_FILES! EQU 1 (
                        echo !LINE!>> "%TEMP_LIST%"
                    )
                )
            )
        )
    )
)

REM Count collectors
set "COLLECTOR_COUNT=0"
for /f "usebackq delims=" %%A in ("%COLLECTOR_LIST%") do (
    set "LINE=%%A"
    if not "!LINE!"=="" (
        echo !LINE! | findstr /C:"collector_" >nul
        if !ERRORLEVEL! EQU 0 (
            set /a COLLECTOR_COUNT+=1
            for %%B in ("!LINE!") do echo       Found: %%~nxB
            echo       Found collector: !LINE! >> "%LOGFILE%"
        )
    )
)

if %COLLECTOR_COUNT% EQU 0 (
    echo       No collector bundles found.
    echo       No collector bundles found. >> "%LOGFILE%"
) else (
    echo       Total collectors found: %COLLECTOR_COUNT%
)

REM =====================================================
REM Download Saved Config (ns.conf from nsconfig folder)
REM Inside collector bundles only
REM =====================================================
echo.
echo [3/5] Downloading saved configurations ^(ns.conf^)...
echo [3/5] Downloading saved configurations... >> "%LOGFILE%"

set "CONFIG_COUNT=0"
for /f "usebackq delims=" %%C in ("%COLLECTOR_LIST%") do (
    set "COLLECTOR_PATH=%%C"
    
    REM Skip empty lines and filter noise
    if not "!COLLECTOR_PATH!"=="" (
        echo !COLLECTOR_PATH! | findstr /C:"collector_" >nul
        if !ERRORLEVEL! EQU 0 (
            for %%D in ("!COLLECTOR_PATH!") do set "COLLECTOR_NAME=%%~nxD"
            
            echo       Checking !COLLECTOR_NAME!...
            
            set "DEST_FILENAME=ns.conf_!COLLECTOR_NAME!"
            
            if "%SCP_TYPE%"=="PSCP" (
                pscp.exe -batch -pw "%PASSWORD%" -q "%USERNAME%@%REMOTE_SERVER%:!COLLECTOR_PATH!/nsconfig/ns.conf" "%DEST%\configs\!DEST_FILENAME!" >nul 2>&1
            ) else (
                scp -o StrictHostKeyChecking=no -q "%USERNAME%@%REMOTE_SERVER%:!COLLECTOR_PATH!/nsconfig/ns.conf" "%DEST%\configs\!DEST_FILENAME!" 2>nul
            )
            
            if exist "%DEST%\configs\!DEST_FILENAME!" (
                echo         [OK] !DEST_FILENAME!
                echo         [OK] Downloaded: !DEST_FILENAME! >> "%LOGFILE%"
                set /a CONFIG_COUNT+=1
            ) else (
                echo         [--] ns.conf not found
            )
        )
    )
)

echo       Saved configs downloaded: %CONFIG_COUNT%

REM =====================================================
REM Download Running Config (ns_running_config.conf)
REM Inside collector bundles only
REM =====================================================
echo.
echo [4/5] Downloading running configurations...
echo [4/5] Downloading running configurations... >> "%LOGFILE%"

set "RUNNING_COUNT=0"
for /f "usebackq delims=" %%C in ("%COLLECTOR_LIST%") do (
    set "COLLECTOR_PATH=%%C"
    
    REM Skip empty lines and filter noise
    if not "!COLLECTOR_PATH!"=="" (
        echo !COLLECTOR_PATH! | findstr /C:"collector_" >nul
        if !ERRORLEVEL! EQU 0 (
            for %%D in ("!COLLECTOR_PATH!") do set "COLLECTOR_NAME=%%~nxD"
            
            echo       Checking !COLLECTOR_NAME!...
            
            set "DEST_FILENAME=ns_running_config.conf_!COLLECTOR_NAME!"
            
            if "%SCP_TYPE%"=="PSCP" (
                pscp.exe -batch -pw "%PASSWORD%" -q "%USERNAME%@%REMOTE_SERVER%:!COLLECTOR_PATH!/shell/ns_running_config.conf" "%DEST%\configs\!DEST_FILENAME!" >nul 2>&1
            ) else (
                scp -o StrictHostKeyChecking=no -q "%USERNAME%@%REMOTE_SERVER%:!COLLECTOR_PATH!/shell/ns_running_config.conf" "%DEST%\configs\!DEST_FILENAME!" 2>nul
            )
            
            if exist "%DEST%\configs\!DEST_FILENAME!" (
                echo         [OK] !DEST_FILENAME!
                echo         [OK] Downloaded: !DEST_FILENAME! >> "%LOGFILE%"
                set /a RUNNING_COUNT+=1
            ) else (
                echo         [--] Running config not found
            )
        )
    )
)

echo       Running configs downloaded: %RUNNING_COUNT%

REM =====================================================
REM Download Capture and Media Files
REM OUTSIDE collector bundles only (case root level)
REM Using pre-fetched file list
REM =====================================================
echo.
echo [5/5] Downloading files from case directory...
echo [5/5] Downloading files from case root... >> "%LOGFILE%"

set "CAPTURE_COUNT=0"
set "MEDIA_COUNT=0"
set "OTHER_COUNT=0"

for /f "usebackq delims=" %%F in ("%TEMP_LIST%") do (
    set "REMOTE_FILE=%%F"
    
    REM Skip empty lines and filter noise
    if not "!REMOTE_FILE!"=="" (
        echo !REMOTE_FILE! | findstr /C:"Keyboard" /C:"prompts" /C:"authentication" >nul
        if !ERRORLEVEL! NEQ 0 (
            for %%G in ("!REMOTE_FILE!") do (
                set "FILENAME=%%~nxG"
                set "FILEEXT=%%~xG"
            )
            
            if not "!FILENAME!"=="" (
                REM Determine file type and destination
                set "FILE_DEST="
                set "FILE_TYPE="
                
                REM Check if capture file
                echo !FILEEXT! | findstr /I /C:".cap" /C:".pcap" /C:".pcapng" /C:".sslkeys" /C:".keys" /C:".har" >nul
                if !ERRORLEVEL! EQU 0 (
                    set "FILE_DEST=%DEST%\captures\!FILENAME!"
                    set "FILE_TYPE=capture"
                )
                
                REM Check if image file
                echo !FILEEXT! | findstr /I /C:".jpg" /C:".jpeg" /C:".png" /C:".gif" /C:".bmp" /C:".webp" /C:".tiff" /C:".svg" >nul
                if !ERRORLEVEL! EQU 0 (
                    set "FILE_DEST=%DEST%\media\!FILENAME!"
                    set "FILE_TYPE=media"
                )
                
                REM Check if other file (txt, log, xml, json, csv)
                echo !FILEEXT! | findstr /I /C:".txt" /C:".log" /C:".xml" /C:".json" /C:".csv" >nul
                if !ERRORLEVEL! EQU 0 (
                    set "FILE_DEST=%DEST%\!FILENAME!"
                    set "FILE_TYPE=other"
                )
                
                REM Download the file
                if defined FILE_DEST (
                    echo       Downloading: !FILENAME!
                    
                    if "%SCP_TYPE%"=="PSCP" (
                        pscp.exe -batch -pw "%PASSWORD%" -q "%USERNAME%@%REMOTE_SERVER%:!REMOTE_FILE!" "!FILE_DEST!" >nul 2>&1
                    ) else (
                        scp -o StrictHostKeyChecking=no -q "%USERNAME%@%REMOTE_SERVER%:!REMOTE_FILE!" "!FILE_DEST!" 2>nul
                    )
                    
                    if exist "!FILE_DEST!" (
                        echo         [OK] !FILENAME!
                        echo         [OK] Downloaded !FILE_TYPE!: !FILENAME! >> "%LOGFILE%"
                        
                        if "!FILE_TYPE!"=="capture" set /a CAPTURE_COUNT+=1
                        if "!FILE_TYPE!"=="media" set /a MEDIA_COUNT+=1
                        if "!FILE_TYPE!"=="other" set /a OTHER_COUNT+=1
                    )
                )
            )
        )
    )
)

echo.
echo       Capture files downloaded: %CAPTURE_COUNT%
echo       Media files downloaded  : %MEDIA_COUNT%
echo       Other files downloaded  : %OTHER_COUNT%

REM =====================================================
REM Cleanup and Summary
REM =====================================================
del "%TEMP_LIST%" 2>nul
del "%TEMP_FILTERED%" 2>nul
del "%COLLECTOR_LIST%" 2>nul
del "%ALL_FILES_LIST%" 2>nul

REM Remove empty directories
for %%D in (configs captures media) do (
    dir /b "%DEST%\%%D" 2>nul | findstr "." >nul || rmdir "%DEST%\%%D" 2>nul
)

REM Refresh icon cache to show new icons
echo.
echo  [ICONS] Refreshing icon cache...
ie4uinit.exe -show >nul 2>&1

echo.
echo  ============================================================
echo   ^|^> Download Complete - Summary
echo  ============================================================
echo.
echo   STORAGE INFORMATION:
echo     Storage Drive            : %STORAGE_DRIVE%
echo     Drive Status             : %DRIVE_STATUS%
echo     Base Path                : %STORAGE_BASE%
echo.
echo   FROM COLLECTOR BUNDLES:
echo     Collector bundles found  : %COLLECTOR_COUNT%
echo     Saved configs ^(ns.conf^)  : %CONFIG_COUNT%
echo     Running configs          : %RUNNING_COUNT%
echo.
echo   FROM CASE ROOT DIRECTORY:
echo     Capture files ^(pcap,har^) : %CAPTURE_COUNT%
echo     Media files ^(images^)     : %MEDIA_COUNT%
echo     Other files ^(xml,txt^)    : %OTHER_COUNT%
echo.
echo   Destination : %DEST%
echo   Log file    : %LOGFILE%
echo.
echo   FOLDER ICONS:
echo     SJLNT folder and case folder have custom icons.
echo     If icons don't appear, press F5 to refresh Explorer.
echo.

REM Write summary to log
echo. >> "%LOGFILE%"
echo ============================================================ >> "%LOGFILE%"
echo  SUMMARY >> "%LOGFILE%"
echo ============================================================ >> "%LOGFILE%"
echo  STORAGE INFORMATION: >> "%LOGFILE%"
echo    Storage Drive            : %STORAGE_DRIVE% >> "%LOGFILE%"
echo    Drive Status             : %DRIVE_STATUS% >> "%LOGFILE%"
echo. >> "%LOGFILE%"
echo  FROM COLLECTOR BUNDLES: >> "%LOGFILE%"
echo    Collector bundles found  : %COLLECTOR_COUNT% >> "%LOGFILE%"
echo    Saved configs            : %CONFIG_COUNT% >> "%LOGFILE%"
echo    Running configs          : %RUNNING_COUNT% >> "%LOGFILE%"
echo. >> "%LOGFILE%"
echo  FROM CASE ROOT DIRECTORY: >> "%LOGFILE%"
echo    Capture files            : %CAPTURE_COUNT% >> "%LOGFILE%"
echo    Media files              : %MEDIA_COUNT% >> "%LOGFILE%"
echo    Other files              : %OTHER_COUNT% >> "%LOGFILE%"
echo ============================================================ >> "%LOGFILE%"

echo  ============================================================
echo.

REM Open destination folder
choice /C YN /M "Open destination folder now?"
if !ERRORLEVEL! EQU 1 explorer "%DEST%"

endlocal
exit /b 0