Install Mac version of UT2004 from PC discs and Mac demo

Main.hs

module Main where

{-
Utility to automatically install and patch the Mac OS X version
of Unreal Tournament 2004 from the Mac demo and patcher
and PC installation disc images and CD key.
Can automatically download the demo and patcher.
see: http://lowetechlabs.com/UT2004-WIN2OSX/
see: http://forums.macnn.com/showthread.php?t=292620

System software dependencies:
    Mac OS X 10.4
Haskell package dependencies:
    HaXml, System.FilePath
External program dependencies:
    hdiutil, defaults, plutil, xmllint, curl, bunzip2
-}

import Control.Concurrent
import Control.Exception (evaluate)
import Control.Monad
import Data.Char
import Data.Maybe
import Network.URI hiding (path)
import System.Console.GetOpt
import System.Environment
import System.Directory
import System.Exit
import System.FilePath
import System.IO
import System.Process
import Text.Regex
import Text.XML.HaXml
import Text.XML.HaXml.Parse
import Text.XML.HaXml.Pretty


installFolder = "/Applications"

launcherFolder = "Contents/MacOS"
utAppName = "Unreal Tournament 2004"
utAppID = "com.epicgames.ut2004"
bundleInfo = "Contents/Info.plist"
cdkeyFile = "System/cdkey"

cdkeyRegex :: Regex
cdkeyRegex = mkRegex
    "[A-Z0-9]{5}-[A-Z0-9]{5}-[A-Z0-9]{5}-[A-Z0-9]{5}"

patchProgramName = "mojopatch"
patchFileName = "default.mojopatch"

pcSpecificExtensions = [".dll", ".exe"]
pcSpecificPathComponents = ["DirectX9"]

uccPath = "System/ucc-bin"
utArchiveExt = ".uz2"
userDemoSettingsFolder
    = "Library/Application Support" </> (utAppName ++ " Demo")
userDemoDecompressionFolder
    = userDemoSettingsFolder </> "System"

utAppPath = installFolder </> (utAppName ++ ".app")

-- see: http://www.unrealtournament.com/ut2004/downloads.html
demoURL :: URI
demoURL = fromJust . parseURI $
    "http://treefort.icculus.org/ut2004/UT2004-MAC-Demo3334.dmg.bz2"

patchPageURL :: URI
patchPageURL = fromJust . parseURI $
    "http://www.apple.com/downloads/macosx/games/demos_updates/unrealtournament2004.html"

usageHeader
     = "UT2004 Install Helper\n"
    ++ "Install the Mac version of UT2004 using your PC discs and CD key.\n"
    ++ "Usage: ut2k4install [OPTION ...] [install-image ...]"

-- entry point

optionDescrs :: [OptDescr InstallerOption]
optionDescrs = [
    Option ['h', '?'] ["help"]        (NoArg    (ShowHelp, ""))
        "Print this usage message and exit.",
    Option ['c'] ["cdkey"]            (ReqArg ((,) CDKey) "CDKEY")
        "The CD key to install, entered exactly as on the package.",
    Option ['d'] ["demo-image"]        (ReqArg    ((,) DemoImage) "DMG")
        "[optional] Path to local copy of the disc image of the demo.",
    Option ['p'] ["patch-image"]    (ReqArg    ((,) PatchImage) "DMG")
        "[optional] Path to local copy of the disc image of a patch.",
    Option ['i'] ["install-image"]    (ReqArg ((,) InstallImage) "ISO")
        "Path to disc image of one of the install CDs or DVDs."
    ]

type InstallerOption = (InstallerOptionName, String)
data InstallerOptionName
    = ShowHelp
    | CDKey
    | DemoImage
    | PatchImage
    | InstallImage
    deriving (Eq, Ord, Show)

main :: IO ()
main = do
    -- parse and validate command-line options
    hSetBuffering stdout NoBuffering
    hSetBuffering stderr NoBuffering
    args <- getArgs
    let
        parseInstallOptions
            = getOpt (ReturnInOrder ((,) InstallImage)) optionDescrs
        (options, _, errors) = parseInstallOptions args
        optionPresent opt = isJust $ lookup opt options
        printUsage = length options == 0 || optionPresent ShowHelp
        optParseError = length errors > 0
        noCDKey = not $ optionPresent CDKey
        cdkey = fromJust $ lookup CDKey options
        badCDKey = not noCDKey && isNothing (matchRegex cdkeyRegex cdkey)
        noInstallImages = not $ optionPresent InstallImage
        errorExists = or [
            printUsage,
            optParseError,
            noCDKey,
            badCDKey,
            noInstallImages
            ]
    if printUsage || optParseError
        then putStrLn (usageInfo usageHeader optionDescrs)
        else return ()
    if optParseError
        then mapM_ putStrLn errors
        else return ()
    if noCDKey
        then putStrLn "Installation requires a CD key."
        else return ()
    if badCDKey
        then putStrLn $
               "The CD key must be in the format XXXXX-XXXXX-XXXXX-XXXXX."
        else return ()
    if noInstallImages
        then putStrLn $
               "Installation requires at least one game disc image."
            ++ " Installing from physical media or a PC UT2004 installation"
            ++ " is not yet supported."
        else return ()
    if errorExists
        then exitFailure
        else do
            installWith options
            exitWith ExitSuccess

installWith :: [InstallerOption] -> IO ()
installWith options = do
    case (lookup DemoImage options) of
        (Just demoImagePath)    -> do
            demoImage <- makeRelativeToCurrentDirectory demoImagePath
            task "Installing demo" (installDemo demoImage)
        Nothing                -> do
            demoArchive <- task "Downloading demo" (fetchTemp demoURL)
            task "Decompressing demo" (exec_ "bunzip2" [demoArchive])
            let demoImage = dropExtension demoArchive
            task "Installing demo" (installDemo demoImage)
            task "Deleting demo installation files" (removeFile demoImage)
            
    task "Altering app bundle" fixBundle
    
    let
        installImagePaths = map snd $ filter ((InstallImage ==) .  fst) options
        n = length installImagePaths
        discTaskLabels = ["Copying files from data disc "
                      ++ show i ++ " of " ++ show n | i <- [1..n]]
    installImages <- mapM makeRelativeToCurrentDirectory installImagePaths
    zipWithM_ task discTaskLabels (map copyDataDisc installImages)
    
    let cdkey = fromJust $ lookup CDKey options
    task "Writing CD key" $ writeFile (utAppPath </> cdkeyFile) cdkey
    
    case (lookup PatchImage options) of
        (Just patchImagePath)    -> do
            patchImage <- makeRelativeToCurrentDirectory patchImagePath
            task "Installing patch" (runPatcher patchImage)
        Nothing                -> do
            patchArchive <- task "Downloading latest patch" downloadPatch
            task "Decompressing demo" (exec_ "bunzip2" [patchArchive])
            let patchImage = dropExtension patchArchive
            task "Installing patch" (runPatcher patchImage)
            task "Deleting patch installation files" (removeFile patchImage)
    
    task "Deleting remaining demo files" $ do
        home <- getHomeDirectory
        let settingsFolder = home </> userDemoSettingsFolder
        removeDirectoryRecursive settingsFolder
    
    putStrLn "Installation complete."
    where
        task message action = do
            putStr (message ++ "... ")
            a <- action
            putStrLn "done."
            return a

-- major installation actions

{-
Copy the UT2004 Mac demo to the installation location,
making sure to preserve symlinks, such as the one from the
app bundle launcher to the actual UT2004 executable.
-}
installDemo :: FilePath -> IO ()
installDemo demoImage = do
    -- Install the demo and change its name to the retail version's.
    demoVolume <- mount demoImage
    demoApp <- liftM head $ listDir demoVolume
    systemCopy
        (demoVolume </> demoApp)
        (utAppPath)
    unmount demoVolume

{-
Change the bundle identifier and other strings
so that the app bundle will be recognized by the patcher.
-}
fixBundle :: IO ()
fixBundle = do
    -- Rename the Mac OS launcher.
    demoLauncher <- liftM head $ listDir
        (installFolder </> (utAppName ++ ".app") </> launcherFolder)
    renameFile
        (utAppPath </> launcherFolder </> demoLauncher)
        (utAppPath </> launcherFolder </> utAppName)
    -- Change some strings in bundle info plist to retail version values.
    let infoPlist = utAppPath </> bundleInfo
    plistSet infoPlist "CFBundleExecutable" utAppName
    plistSet infoPlist "CFBundleName" utAppName
    plistSet infoPlist "CFBundleIdentifier" utAppID
    -- Munge info plist into a mojopatch-compatible format.
    exec "plutil" ["-convert", "xml1", infoPlist]
    reorderPlistForMojopatch infoPlist
    exec_ "xmllint" ["--format", "--output", infoPlist, infoPlist]

{-
Mojopatch is the patch program used by UT2004 on the Mac.
see: http://icculus.org/mojopatch/
It's really stupid for a Mac OS X app.
First, it gets application versions from their Info.plist files,
but it uses its own plist routines, so it doesn't understand plists
in binary format. For XML plists, it uses an internal XML parser,
which is a complete joke that will break if it sees a dict or array
value before it finds the product version.
All I can say is, thanks, Mojopatch author, for making it open source.
I never would have figured out how to work around that mess otherwise.
Apple's plutil and defaults utilities will order keys alphabetically.
This function puts keys with simple values (strings, booleans, etc.)
at the beginning of the plist's root dictionary and other keys after that,
so mojopatch doesn't get confused and die.
-}
reorderPlistForMojopatch :: FilePath -> IO ()
reorderPlistForMojopatch plistFile = do
    plistText <- readFile plistFile
    readMVar <- newEmptyMVar
    forkIO (evaluate (length plistText) >> putMVar readMVar ())
    takeMVar readMVar
    let doc = xmlParse plistFile plistText
    writeFile plistFile
        $ render . document
        $ updateRoot doc . head
        $ rewritePlist $ getRoot doc
    where
        updateRoot :: Document -> Content -> Document
        updateRoot (Document prolog st _ misc) (CElem root)
            = Document prolog st root misc
        
        simpleVTags = ["string", "real", "integer",
                       "date", "true", "false", "data"]
        otherVTags = ["dict", "array"]
        
        entries :: [String] -> CFilter
        entries vTags = sep . (tagged (elm `o` children `o` tag "dict"))
            where
                sep (("key", k):(vTag, v):es)
                    | vTag `elem` vTags = k:v:(sep es)
                    | otherwise = sep es
                sep ((n, t):es) = error $
                       "Couldn't prepare app bundle for patching."
                    ++ " Malformed plist file, not expecting tag " ++ show n
                sep [] = []
        
        simpleTagsInFront :: CFilter
        simpleTagsInFront
            =   pip (union (entries simpleVTags) (entries otherVTags))
            `o` tag "dict"
        
        rewritePlist :: CFilter
        rewritePlist = chip simpleTagsInFront `o` tag "plist"
        
        {-
        Process In Place: sort of like HaXml's chip.
        Applies a filter to an element and replaces its children
        with the results of that filter.
        -}
        pip :: CFilter -> CFilter
        pip f c@(CElem (Elem n as _)) = [ CElem (Elem n as (f c)) ]
        pip f c = [c]

-- HaXml utility functions
getRoot :: Document -> Content
getRoot (Document _ _ root _) = CElem root

getChildText :: Content -> String
getChildText = concatMap getText . (txt `o` children)

getText :: Content -> String
getText (CString _ string) = string
getText _ = error $
    "Error parsing XML. Can't get text from non-text content."
        
getAttr :: String -> Content -> String
getAttr name (CElem (Elem _ attrs _)) =
    getValue . fromJust . lookup name $ attrs
    where
        getValue (AttValue pieces) = concatMap valuePiece pieces
        valuePiece (Left str) = str
        valuePiece (Right (RefChar c)) = [chr c]
        -- Why doesn't HaXml resolve entities automatically?
        -- Standard XML 1.0 named entities:
        valuePiece (Right (RefEntity "amp")) = "&"
        valuePiece (Right (RefEntity "lt")) = "<"
        valuePiece (Right (RefEntity "gt")) = ">"
        valuePiece (Right (RefEntity "apos")) = "'"
        valuePiece (Right (RefEntity "quot")) = "\""
        valuePiece (Right (RefEntity name)) = error $
            "Error parsing XML. Unknown named entity: " ++ show name

{-
Mount an install disc image from the PC version,
copy data files into the Mac OS X application bundle,
uncompressing any compressed files,and unmount the image
when done. Ignores files in the root folder of the disc
because they're all PC-specific.
-}
copyDataDisc :: FilePath -> IO ()
copyDataDisc image = do
    volume <- mount image
    dataFolders <- filterM doesDirectoryExist
                =<< (liftM (map (volume </>)) $ listDir volume)
    zipWithM_ (copyChecked isNotPCSpecific)
        dataFolders
        (map ((utAppPath </>) . takeFileName) dataFolders)
    unmount volume

{-
Download the latest UT2004 Mac patch from Apple.
-}
downloadPatch :: IO FilePath
downloadPatch = do
    -- Query the Apple game info page for UT to find the latest patch.
    patchPageFile <- fetchTemp patchPageURL
    -- HaXml can't parse real HTML even with its correcting parser.
    text <- exec "xmllint" ["--html", "--xmlout", patchPageFile]
    removeFile patchPageFile
    doc <- case (xmlParse' patchPageFile text) of
        (Left errorMsg)    -> fail $
               "Error downloading latest patch."
            ++ " Could not parse patch info page: "
            ++ errorMsg
        (Right doc)        -> return doc
    let
        root = getRoot doc
        patchURL = fromJust . parseURI . selectPatchLink $ root
    patchImage <- fetchTemp patchURL
    return patchImage
    where
        selectPatchLink = getAttr "href" . head . (deep $ path
            [tag "a", attrval (mkAttr "class" "download")])
                        
        mkAttr :: Name -> String -> Attribute
        mkAttr name value = (name, AttValue [(Left value)])

{-
Run the UT patch application.
This function assumes the patch app is mojopatch.
Patch 3369-2 is required for Intel Mac support.
-}
runPatcher :: FilePath -> IO ()
runPatcher patchImage = do
    patchVolume <- mount patchImage
    patchApp <- liftM head $ listDir patchVolume
    let
        patchAppPath = patchVolume </> patchApp
        patchProgram = patchAppPath </> launcherFolder </> patchProgramName
        patchFile = patchAppPath </> patchFileName
    exec patchProgram ["--ui", "stdio", "--quietonsuccess", patchFile]
    unmount patchVolume

-- support functions

-- Run an external program and return its output.
exec :: String -> [String] -> IO String
exec programName args = do
    programPath <- findExecutable programName
    (_, out, _, pid) <- runInteractiveProcess programName args Nothing Nothing
    output <- hGetContents out
    outMVar <- newEmptyMVar
    forkIO (evaluate (length output) >> putMVar outMVar ())
    takeMVar outMVar
    exitCode <- waitForProcess pid
    case exitCode of
        ExitSuccess                -> return output
        ExitFailure failCode    -> fail $
            "Error executing external program." ++
            " Program " ++ (show programName) ++
            " with args " ++ (show args) ++
            " failed with exit code " ++ (show failCode)

-- Run an external program, ignoring its output.
exec_ :: String -> [String] -> IO ()
exec_ programName args = do
    programPath <- findExecutable programName
    (_, _, _, pid) <- runInteractiveProcess programName args Nothing Nothing
    exitCode <- waitForProcess pid
    case exitCode of
        ExitSuccess                -> return ()
        ExitFailure failCode    -> fail $
            "Error executing external program." ++
            " Program " ++ (show programName) ++
            " with args " ++ (show args) ++
            " failed with exit code " ++ (show failCode)

-- Mount a disk image. Returns the path of the mounted volume.
mount :: FilePath -> IO (FilePath)
mount image = do
    mountList <- exec "hdiutil" ["mount", "-plist", image]
    let
        doc = xmlParse image mountList
        kvList = foldKeyValuePairs
            $ multi (keysAndValues `o` tag "dict" `o` children)
            $ getRoot doc
        volumes = multilookup "mount-point" kvList
    case (length volumes) of
        1 -> return (head volumes)
        0 -> fail $
            "Error mounting disk image. No volumes found: " ++ image
        _ -> fail $
            "Error mounting disk image. Multiple volumes found, " ++
            "but we don't know which one is the right one: " ++ image
    where
        keysAndValues :: CFilter
        keysAndValues dict = mkPairs $ elm `o` children $ dict
            where
                mkPairs (k:v:xs)
                    | null (tag "key" k) = (mkPairs xs)
                    | otherwise = k:v:(mkPairs xs)
                mkPairs _ = []
        
        foldKeyValuePairs (k:v:xs)
            = (getChildText k, getChildText v):(foldKeyValuePairs xs)
        foldKeyValuePairs _ = []
        
        multilookup key = map snd . filter ((key ==) . fst)

-- Unmount a mounted volume and all associated partitions.
unmount :: FilePath -> IO ()
unmount volume = exec_ "hdiutil" ["detach", volume]

{-
Set a value in a plist file using Apple's defaults utility.
According to the defaults man page, this functionality may be
moved out of defaults in the next release of Mac OS X.
A side effect of this use of defaults is that it converts
the plist it's operating on to the binary plist format.
-}
plistSet :: FilePath -> String -> String -> IO ()
plistSet plist key value
    = exec_ "defaults" ["write", dropExtensions plist, key, value]

{-
Get the filenames in a directory, ignoring hidden files
and the navigation entries "." and "..".
-} 
listDir :: FilePath -> IO [FilePath]
listDir path = do
    childFilenames <- getDirectoryContents path
    return (filter isVisible $ childFilenames)
    where isVisible file = head file /= '.'

-- Copy recursively, preserving symlinks and resource forks.
systemCopy :: FilePath -> FilePath -> IO ()
systemCopy src dst = exec_ "cp" ["-R", src, dst]

{-
Recursive copy routine for merging folder structures.
Checks paths to see if they should be copied.
It will not overwrite directories (which would remove all files
in the destination directory); instead, only files are actually copied,
and directories are created as needed.
-}
copyChecked :: (FilePath -> Bool) -> FilePath -> FilePath -> IO ()
copyChecked test src dst = do
    let
        isValid            = test            src
        isCompressed     = isUTArchive    src
    isFile        <- doesFileExist        src
    isFolder    <- doesDirectoryExist    src
    case (isValid, isCompressed, isFile, isFolder) of
        -- regular file
        (True, False, True, False)    -> do
            createDirectoryIfMissing True (takeDirectory dst)
            copyFile src dst
        -- regular folder
        (True, False, False, True)    -> do
            createDirectoryIfMissing True dst
            children <- listDir src
            zipWithM_ (copyChecked test)
                (map (src </>) children)
                (map (dst </>) children)
        -- compressed file
        (True, True, True, False)    -> do
            createDirectoryIfMissing True (takeDirectory dst)
            decompress src (dropExtension dst)
        -- something we should ignore
        (False, _, _, _)            -> return ()
        _                            -> fail $
            (show src) ++ " does not exist, or is not a file or folder"

-- Pass only files that aren't specific to the PC version of UT 2004.
isNotPCSpecific :: FilePath -> Bool
isNotPCSpecific path = not $
       (map toLower (takeExtension path)) `elem` pcSpecificExtensions
    || or (map (`elem` (splitDirectories path)) pcSpecificPathComponents)

isUTArchive :: FilePath -> Bool
isUTArchive path = utArchiveExt == map toLower (takeExtension path)

{-
Expand an Unreal archive using the "UnrealOS execution environment".
see: http://wiki.beyondunreal.com/wiki/Ucc

For some incomprehensible reason, instead of being expanded to the
working folder or the Unreal install folder, archives are expanded
using the System folder in the per-user Unreal settings folder
as the prefix for the input filename with the compressed suffixed removed.

Since the files that tell it to be the full version haven't been copied
yet, this function actually decompresses to the demo's settings folder.
Must be run between installing demo and overwriting demo files
or removing demo per-user settings folder.

This is dumb so here's an example:
/Volumes/UTDisc1/Maps/ExampleMap.ut2.uz2 ->
~/Library/Application Support/Unreal Tournament 2004 Demo/System
    </> /Volumes/UTDisc1/Maps/ExampleMap.ut2
-}
decompress :: FilePath -> FilePath -> IO ()
decompress src dst = do
    home <- getHomeDirectory
    let
        decompressionFolder = home </> userDemoDecompressionFolder
        ucc = utAppPath </> uccPath
        decompressedFilePath
            = decompressionFolder </> (makeRelative "/" (dropExtension src))
    exec ucc ["decompress", src]
    renameFile decompressedFilePath dst
    return ()

-- Retrieve a file from a URL and put it in the system temporary file folder.
fetchTemp :: URI -> IO FilePath
fetchTemp uri = do
    file <- getTempFileFor (filenameFromURI uri)
    fetch uri file
    return file
    where
        -- Return the name of the file referenced by a URL.
        filenameFromURI :: URI -> String
        filenameFromURI
            = reverse . takeWhile (/= '/') . reverse . uriPath
        
        -- Retrieve a file from a URL and write it to the specified path.
        fetch :: URI -> FilePath -> IO ()
        fetch uri file = exec_ "curl" ["--silent", "--location",
            "--output", file, uriToString id uri ""]
        
        -- Return an unused file name in the system temporary file folder.
        getTempFileFor :: String -> IO FilePath
        getTempFileFor filename = do
            tempFolder <- getTemporaryDirectory
            (file, handle) <- openTempFile tempFolder filename
            hClose handle
            removeFile file
            return file