@echo off
setlocal

set ERROR=0

if /i "%1"==""      goto usage
if /i "%1"=="/?"    goto usage
if /i "%1"=="-?"    goto usage
if /i "%1"=="/h"    goto usage
if /i "%1"=="-h"    goto usage
if /i "%1"=="/help" goto usage
if /i "%1"=="-help" goto usage

if /i not "%3"=="" (
    echo ERROR: Too many arguments
    goto usage
)

if /i "%2"==""  goto validSetting
if /i "%2"=="1" goto validSetting
if /i "%2"=="2" goto validSetting
if /i "%2"=="3" goto validSetting
if /i "%2"=="4" goto validSetting

echo ERROR: Invalid SAN policy specified: %2
goto usage

:validSetting

REM
REM Find the system hive
REM

set BASEDIR=%~f1

set SYSTEMHIVE=%BASEDIR%\windows\system32\config\SYSTEM
if exist %SYSTEMHIVE% (
    goto performOperation
)

set SYSTEMHIVE=%BASEDIR%\system32\config\SYSTEM
if exist %SYSTEMHIVE% (
    goto performOperation
)

echo ERROR: Invalid Windows PE image: %BASEDIR%
goto usage

:performOperation

REM
REM Load the system hive
REM

set HIVESTRING=%SYSTEMHIVE:\=-%
set HIVESTRING=%HIVESTRING:/=-%

reg load HKLM\%HIVESTRING% "%SYSTEMHIVE%" > NUL
if errorlevel 1 (
    echo ERROR: Could not mount the Windows PE SYSTEM hive: %ERRORLEVEL%
    goto fail
)

if /i "%2"=="" goto getCurrentSetting
 
reg add HKLM\%HIVESTRING%\ControlSet001\Services\partmgr\Parameters /v SanPolicy /t REG_DWORD /d %2 /f > NUL
if errorlevel 1 (
    echo ERROR: Could not set the SAN policy: %ERRORLEVEL%
    set ERROR=%ERRORLEVEL%
) else (
    echo Successfully changed the SAN policy.
)

:getCurrentSetting

echo Current setting:
reg query HKLM\%HIVESTRING%\ControlSet001\Services\partmgr\Parameters /v SanPolicy | findstr /i SanPolicy
if errorlevel 1 (
    echo ERROR: Could not query the SAN policy: %ERRORLEVEL%
    set ERROR=%ERRORLEVEL%
)

set UNLOAD=
:unload
set UNLOAD=%UNLOAD%.
reg unload HKLM\%HIVESTRING% 2>NUL > NUL

if errorlevel 1 (

    echo Attempting to unload the Windows PE SYSTEM hive%UNLOAD%

    if "%UNLOAD%" GEQ ".........." (
        echo.
        reg unload HKLM\foo > NUL

        if errorlevel 1 (
            echo ERROR: Could not unload the Windows PE SYSTEM hive: %ERRORLEVEL%
            echo.                                                                                
            echo You will need to manually unload the hive at "HKLM\%HIVESTRING%".  The policy
            echo may not persist in the WinPE image until the hive is successfully unloaded.
            set ERROR=%ERRORLEVEL%
            goto fail
        )

    ) else (
        goto unload
    )
)

if not "0" == "%ERROR%" (
    goto fail
)

:success
set ERROR=0
echo Setsanpolicy completed successfully.
goto cleanup

:fail
if "0" == "%ERROR%" (
    if errorlevel 1 (
        set ERROR=%ERRORLEVEL%
    ) else (
        set ERROR=1
    )
)
echo Setsanpolicy failed.
goto cleanup

:usage
set ERROR=1
echo Usage: setsanpolicy image [ 1 ^| 2 ^| 3 ]
echo.
echo 'image' - the path to a Windows PE image
echo.
echo 1 = Online  all disks (default)
echo 2 = Online  all disks except those on a shared bus
echo 3 = Offline all disks
echo.
echo Omitting a policy number prints the current policy.
echo.
echo Example: setsanpolicy c:\winpe\mount 2
goto cleanup

:cleanup
endlocal & exit /b %ERROR%