diff options
| -rw-r--r-- | src/Environment.hs | 81 | ||||
| -rw-r--r-- | src/Registry.hs | 218 |
2 files changed, 189 insertions, 110 deletions
diff --git a/src/Environment.hs b/src/Environment.hs index 2b53258..5a3978e 100644 --- a/src/Environment.hs +++ b/src/Environment.hs @@ -5,9 +5,9 @@ -} module Environment - ( saveToRegistry + ( queryFromRegistry + , saveToRegistry , saveToRegistryWithPrompt - , queryFromRegistry , wipeFromRegistry , wipeFromRegistryWithPrompt , getEnv @@ -16,12 +16,15 @@ module Environment , RegistryBasedEnvironment(..) ) where -import Control.Monad (liftM, when) -import Data.List (intercalate) -import Data.List.Split (splitOn) -import Data.Maybe (fromMaybe) +import Control.Monad (liftM, when) +import Data.List (intercalate) +import Data.List.Split (splitOn) +import Data.Maybe (fromMaybe) import qualified System.Environment (lookupEnv) -import System.IO.Error (catchIOError, isDoesNotExistError) +import System.IO.Error (catchIOError, isDoesNotExistError) + +import qualified Graphics.Win32.Window as WinAPI +import qualified System.Win32.Types as WinAPI import qualified Registry import qualified Utils (promptToContinue) @@ -31,22 +34,45 @@ data RegistryBasedEnvironment | AllUsersEnvironment deriving (Eq, Show) -registrySubKeyPath :: RegistryBasedEnvironment -> String -registrySubKeyPath CurrentUserEnvironment = "Environment" -registrySubKeyPath AllUsersEnvironment = "SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment" +subKeyPath :: RegistryBasedEnvironment -> Registry.KeyPath +subKeyPath CurrentUserEnvironment = + Registry.keyPathFromString "Environment" +subKeyPath AllUsersEnvironment = + Registry.keyPathFromString "SYSTEM\\CurrentControlSet\\Control\\Session Manager\\Environment" + +rootKey :: RegistryBasedEnvironment -> Registry.RootKey +rootKey CurrentUserEnvironment = Registry.CurrentUser +rootKey AllUsersEnvironment = Registry.LocalMachine + +openRootKey :: RegistryBasedEnvironment -> Registry.KeyHandle +openRootKey = Registry.openRootKey . rootKey -registryKey :: RegistryBasedEnvironment -> Registry.KeyHandle -registryKey CurrentUserEnvironment = Registry.hkcu -registryKey AllUsersEnvironment = Registry.hklm +openRegistryKey :: RegistryBasedEnvironment -> IO Registry.KeyHandle +openRegistryKey env = Registry.openSubKey (openRootKey env) (subKeyPath env) + +registryKeyPath :: RegistryBasedEnvironment -> Registry.KeyPath +registryKeyPath env = Registry.keyPathJoin [Registry.rootKeyPath $ rootKey env, subKeyPath env] + +notifyEnvUpdate :: IO () +notifyEnvUpdate = + WinAPI.withTString "Environment" $ \lparamPtr -> do + let wparam = 0 + let lparam = fromIntegral $ WinAPI.castPtrToUINTPtr lparamPtr + _ <- WinAPI.sendMessage allWindows messageCode wparam lparam + return () + where + wM_SETTINGCHANGE = 0x1A + messageCode = wM_SETTINGCHANGE -registryKeyPath :: RegistryBasedEnvironment -> String -registryKeyPath CurrentUserEnvironment = "HKCU\\" ++ registrySubKeyPath CurrentUserEnvironment -registryKeyPath AllUsersEnvironment = "HKLM\\" ++ registrySubKeyPath AllUsersEnvironment + allWindows = WinAPI.castUINTPtrToPtr 0xffff -saveToRegistry :: RegistryBasedEnvironment -> String -> String -> IO () -saveToRegistry env = Registry.setString (registryKey env) (registrySubKeyPath env) +saveToRegistry :: RegistryBasedEnvironment -> Registry.ValueName -> Registry.ValueData -> IO () +saveToRegistry env name value = do + keyHandle <- openRegistryKey env + Registry.setString keyHandle name value + notifyEnvUpdate -saveToRegistryWithPrompt :: RegistryBasedEnvironment -> String -> String -> IO () +saveToRegistryWithPrompt :: RegistryBasedEnvironment -> Registry.ValueName -> Registry.ValueData -> IO () saveToRegistryWithPrompt env name value = do putStrLn $ "Saving variable '" ++ name ++ "' to '" ++ registryKeyPath env ++ "'..." oldValue <- queryFromRegistry env name @@ -55,19 +81,22 @@ saveToRegistryWithPrompt env name value = do agreed <- Utils.promptToContinue when agreed $ saveToRegistry env name value -queryFromRegistry :: RegistryBasedEnvironment -> String -> IO String -queryFromRegistry env name = catchIOError (Registry.getString (registryKey env) (registrySubKeyPath env) name) emptyIfDoesNotExist +queryFromRegistry :: RegistryBasedEnvironment -> Registry.ValueName -> IO Registry.ValueData +queryFromRegistry env name = do + keyHandle <- openRegistryKey env + catchIOError (Registry.getString keyHandle name) emptyIfDoesNotExist where - emptyIfDoesNotExist :: IOError -> IO String emptyIfDoesNotExist e = if isDoesNotExistError e then return "" else ioError e -wipeFromRegistry :: RegistryBasedEnvironment -> String -> IO () -wipeFromRegistry env name = catchIOError (Registry.delValue (registryKey env) (registrySubKeyPath env) name) ignoreIfDoesNotExist +wipeFromRegistry :: RegistryBasedEnvironment -> Registry.ValueName -> IO () +wipeFromRegistry env name = do + keyHandle <- openRegistryKey env + catchIOError (Registry.delValue keyHandle name) ignoreIfDoesNotExist + notifyEnvUpdate where - ignoreIfDoesNotExist :: IOError -> IO () ignoreIfDoesNotExist e = if isDoesNotExistError e then return () else ioError e -wipeFromRegistryWithPrompt :: RegistryBasedEnvironment -> String -> IO () +wipeFromRegistryWithPrompt :: RegistryBasedEnvironment -> Registry.ValueName -> IO () wipeFromRegistryWithPrompt env name = do putStrLn $ "Deleting variable '" ++ name ++ "' from '" ++ registryKeyPath env ++ "'..." agreed <- Utils.promptToContinue diff --git a/src/Registry.hs b/src/Registry.hs index a38ef00..6ca9591 100644 --- a/src/Registry.hs +++ b/src/Registry.hs @@ -5,96 +5,146 @@ -} module Registry - ( KeyHandle + ( KeyPath + , keyPathFromString + , keyPathJoin + , keyPathSplit + + , KeyHandle + , openSubKey + + , RootKey(..) + , openRootKey + , rootKeyPath + + , ValueName , delValue + + , ValueData , getString , setString - , hkcu - , hklm ) where -import Control.Exception (bracket) -import Data.Maybe (fromMaybe) -import Foreign.ForeignPtr (withForeignPtr) +import Data.List (intercalate) +import Data.List.Split (splitOn) +import Foreign.ForeignPtr (withForeignPtr) import Foreign.Marshal.Alloc (alloca, allocaBytes) -import Foreign.Ptr (castPtr, plusPtr) -import Foreign.Storable (peek, poke, sizeOf) -import Graphics.Win32.Window (sendMessage) -import System.IO.Error (mkIOError, doesNotExistErrorType) - -import System.Win32.Types -import System.Win32.Registry - -newtype KeyHandle = KeyHandle HKEY - -getType :: HKEY -> String -> String -> IO (Maybe RegValueType) -getType key subKeyPath valueName = - bracket (regOpenKey key subKeyPath) regCloseKey $ \hKey -> - withForeignPtr hKey $ \p_key -> - withTString valueName $ \p_valueName -> - alloca $ \p_type -> do - ret <- c_RegQueryValueEx p_key p_valueName nullPtr p_type nullPtr nullPtr - case ret of - 0x0 -> do - type_ <- peek p_type - return $ Just type_ - 0x2 -> return Nothing - _ -> failWith "RegQueryValueEx" ret - -getString :: KeyHandle -> String -> String -> IO String -getString (KeyHandle hKey) subKeyPath valueName = - bracket (regOpenKey hKey subKeyPath) regCloseKey $ \hSubKey -> - withForeignPtr hSubKey $ \p_key -> - withTString valueName $ \p_valueName -> +import Foreign.Ptr (castPtr, plusPtr) +import Foreign.Storable (peek, poke, sizeOf) +import System.IO.Error (catchIOError, doesNotExistErrorType, mkIOError, isDoesNotExistError) + +import qualified System.Win32.Types as WinAPI +import qualified System.Win32.Registry as WinAPI + +type KeyName = String +type KeyPath = KeyName + +keyPathSep :: KeyPath +keyPathSep = "\\" + +keyPathFromString :: String -> KeyPath +keyPathFromString = keyPathJoin . keyPathSplit + +keyPathSplit :: KeyPath -> [KeyName] +keyPathSplit = filter (not . null) . splitOn keyPathSep + +keyPathJoin :: [KeyName] -> KeyPath +keyPathJoin = intercalate keyPathSep . filter (not . null) + +type KeyHandle = WinAPI.HKEY + +openSubKey :: KeyHandle -> KeyPath -> IO KeyHandle +openSubKey = WinAPI.regOpenKey + +data RootKey = CurrentUser + | LocalMachine + deriving (Eq, Show) + +rootKeyPath :: RootKey -> KeyName +rootKeyPath CurrentUser = "HKCU" +rootKeyPath LocalMachine = "HKLM" + +openRootKey :: RootKey -> KeyHandle +openRootKey CurrentUser = WinAPI.hKEY_CURRENT_USER +openRootKey LocalMachine = WinAPI.hKEY_LOCAL_MACHINE + +type ValueName = String + +raiseDoesNotExistError :: String -> IO a +raiseDoesNotExistError functionName = + ioError $ mkIOError doesNotExistErrorType functionName Nothing Nothing + +raiseUnknownError :: String -> WinAPI.ErrCode -> IO a +raiseUnknownError functionName exitCode = WinAPI.failWith functionName exitCode + +exitCodeSuccess :: WinAPI.ErrCode +exitCodeSuccess = 0 + +exitCodeFileNotFound :: WinAPI.ErrCode +exitCodeFileNotFound = 0x2 + +exitCodeMoreData :: WinAPI.ErrCode +exitCodeMoreData = 0xea + +raiseError :: String -> WinAPI.ErrCode -> IO a +raiseError functionName ret + | ret == exitCodeFileNotFound = raiseDoesNotExistError functionName + | otherwise = raiseUnknownError functionName ret + +delValue :: KeyHandle -> ValueName -> IO () +delValue keyHandle valueName = + withForeignPtr keyHandle $ \keyPtr -> + WinAPI.withTString valueName $ \valueNamePtr -> do + ret <- WinAPI.c_RegDeleteValue keyPtr valueNamePtr + if ret == exitCodeSuccess + then return () + else raiseError "RegDeleteValue" ret + +type ValueType = WinAPI.RegValueType + +getType :: KeyHandle -> ValueName -> IO ValueType +getType keyHandle valueName = + withForeignPtr keyHandle $ \keyPtr -> + WinAPI.withTString valueName $ \valueNamePtr -> + alloca $ \typePtr -> do + ret <- WinAPI.c_RegQueryValueEx keyPtr valueNamePtr WinAPI.nullPtr typePtr WinAPI.nullPtr WinAPI.nullPtr + if ret == exitCodeSuccess + then peek typePtr + else raiseError "RegQueryValueEx" ret + +type ValueData = String + +getString :: KeyHandle -> ValueName -> IO ValueData +getString keyHandle valueName = + withForeignPtr keyHandle $ \keyPtr -> + WinAPI.withTString valueName $ \valueNamePtr -> alloca $ \dataSizePtr -> do poke dataSizePtr 0 - ret <- c_RegQueryValueEx p_key p_valueName nullPtr nullPtr nullPtr dataSizePtr - case ret of - 0x0 -> do - dataSize <- peek dataSizePtr - let newDataSize = dataSize + fromIntegral (sizeOf (undefined :: TCHAR)) - poke dataSizePtr newDataSize - allocaBytes (fromIntegral newDataSize) $ \dataPtr -> do - poke (castPtr $ plusPtr dataPtr $ fromIntegral dataSize) '\0' - failUnlessSuccess "RegQueryValueEx" $ - c_RegQueryValueEx p_key p_valueName nullPtr nullPtr dataPtr dataSizePtr - peekTString $ castPtr dataPtr - 0x2 -> ioError $ mkIOError doesNotExistErrorType "RegQueryValueEx" Nothing $ Just (subKeyPath ++ "\\" ++ valueName) - _ -> failWith "RegQueryValueEx" ret - -setString :: KeyHandle -> String -> String -> String -> IO () -setString (KeyHandle hKey) subKeyPath valueName valueValue = - bracket (regOpenKey hKey subKeyPath) regCloseKey $ \subKey -> - withTString valueValue $ \p_valueValue -> do - type_ <- getType hKey subKeyPath valueName - regSetValueEx subKey valueName (fromMaybe rEG_SZ type_) p_valueValue $ (length valueValue + 1) * sizeOf (undefined :: TCHAR) - notifyEnvironmentUpdate - -notifyEnvironmentUpdate :: IO () -notifyEnvironmentUpdate = - withTString "Environment" $ \p_lparam -> do - let wparam = 0 - let lparam = fromIntegral $ castPtrToUINTPtr p_lparam - let hwnd = castUINTPtrToPtr 0xffff - _ <- sendMessage hwnd wM_SETTINGCHANGE wparam lparam - return () + ret <- WinAPI.c_RegQueryValueEx keyPtr valueNamePtr WinAPI.nullPtr WinAPI.nullPtr WinAPI.nullPtr dataSizePtr + if ret == exitCodeSuccess + then return "" + else if ret /= exitCodeMoreData + then raiseError "RegQueryValueEx" ret + else getStringTerminated keyPtr valueNamePtr dataSizePtr + where + getStringTerminated keyPtr valueNamePtr dataSizePtr = do + dataSize <- peek dataSizePtr + let newDataSize = dataSize + fromIntegral (sizeOf (undefined :: WinAPI.TCHAR)) + poke dataSizePtr newDataSize + allocaBytes (fromIntegral newDataSize) $ \dataPtr -> do + poke (castPtr $ plusPtr dataPtr $ fromIntegral dataSize) '\0' + ret <- WinAPI.c_RegQueryValueEx keyPtr valueNamePtr WinAPI.nullPtr WinAPI.nullPtr dataPtr dataSizePtr + if ret == exitCodeSuccess + then WinAPI.peekTString $ castPtr dataPtr + else raiseError "RegQueryValueEx" ret + +setString :: KeyHandle -> ValueName -> ValueData -> IO () +setString key name value = + WinAPI.withTString value $ \valuePtr -> do + type_ <- catchIOError (getType key name) stringTypeByDefault + WinAPI.regSetValueEx key name type_ valuePtr valueSize where - wM_SETTINGCHANGE = 0x1A - -delValue :: KeyHandle -> String -> String -> IO () -delValue (KeyHandle hKey) subKeyPath valueName = - bracket (regOpenKey hKey subKeyPath) regCloseKey $ \subKey -> - withForeignPtr subKey $ \subKeyPtr -> - withTString valueName $ \p_valueName -> do - ret <- c_RegDeleteValue subKeyPtr p_valueName - notifyEnvironmentUpdate - case ret of - 0x0 -> return () - 0x2 -> ioError $ mkIOError doesNotExistErrorType "RegQueryValueEx" Nothing $ Just (subKeyPath ++ "\\" ++ valueName) - _ -> failWith "RegDeleteValue" ret - -hkcu :: KeyHandle -hkcu = KeyHandle hKEY_CURRENT_USER - -hklm :: KeyHandle -hklm = KeyHandle hKEY_LOCAL_MACHINE + stringTypeByDefault e = if isDoesNotExistError e + then return WinAPI.rEG_SZ + else ioError e + valueSize = (length value + 1) * sizeOf (undefined :: WinAPI.TCHAR) |