Program.fs

// Learn more about F# at http://fsharp.org
//
//   - backup file list
//     /%appdata%/mbackup/mbackup-default.list
//     /%appdata%/mbackup/user-default.list
//     /%appdata%/mbackup/local.list (optional)
//   - exclude pattern
//     /%appdata%/mbackup/mbackup-default.exclude
//     /%appdata%/mbackup/local.exclude (optional)

module Mbackup.Program

open System
open System.IO
open System.Diagnostics
open System.Text.RegularExpressions
open System.Diagnostics.CodeAnalysis

open Argu

open Mbackup.Lib
open Mbackup.ConfigParser

let ExitSuccess = 0
let ExitBadParam = 1
let ExitTimeout = 2
let ExitIOError = 3
let ExitUserError = 4

let version = Reflection.Assembly.GetEntryAssembly().GetName().Version
let versionStr = version.ToString()

[<SuppressMessage("*", "UnionCasesNames")>]
type CLIArguments =
    | [<AltCommandLine("-n")>] Dry_Run
    | Target of backupTarget: string
    | Remote_User of remoteUser: string
    | [<AltCommandLine("-i")>] Itemize_Changes
    | Node_Name of nodeName: string
    | Ssh_Key of sshKeyFilename: string
    | [<AltCommandLine("-V")>] Version
    interface IArgParserTemplate with
        member s.Usage =
            match s with
            | Dry_Run _ -> "only show what will be done, do not transfer any file"
            | Target _ -> "rsync target, could be local dir in Windows or mingw format or remote ssh dir"
            | Remote_User _ -> "remote linux user to own the backup files"
            | Itemize_Changes _ -> "add -i option to rsync"
            | Node_Name _ -> "local node's name, used in remote logging"
            | Ssh_Key _ -> "ssh private key, used when backup to remote ssh node"
            | Version _ -> "show mbackup version and exit"

let programFilesDirWin = Environment.GetFolderPath(Environment.SpecialFolder.ProgramFiles) |> ensureWinDir
let programFilesDir = toMingwPath programFilesDirWin
let mbackupProgramDirWin = programFilesDirWin + "mbackup\\"
let mbackupProgramDir = toMingwPath mbackupProgramDirWin

let appDataRoamingDir =
    Environment.GetFolderPath(Environment.SpecialFolder.ApplicationData)
    |> toMingwPath
    |> ensureDir

let programDataDirWin = getEnv "PROGRAMDATA" |> ensureWinDir
let programDataDir = toMingwPath programDataDirWin
let appDataLocalDirWin = Environment.GetFolderPath(Environment.SpecialFolder.LocalApplicationData) |> ensureWinDir
let appDataLocalDir = appDataLocalDirWin |> toMingwPath

let mbackupInstallDirWin =
    Environment.GetFolderPath(Environment.SpecialFolder.ProgramFiles)
    |> ensureDir
    |> fun s -> s + "mbackup"

let mbackupInstallDir = mbackupInstallDirWin |> toMingwPath

let userHomeWin =
    getEnvDefault "HOME" (Environment.GetFolderPath(Environment.SpecialFolder.UserProfile)) |> ensureWinDir

let userHome = userHomeWin |> toMingwPath

let userConfigDirWin = programDataDirWin + "mbackup\\"
let userConfigDir = programDataDir + "mbackup/"
let runtimeDirWin = appDataLocalDirWin + "mbackup\\"
let runtimeDir = appDataLocalDir + "mbackup/"

let mbackupConfigFile = userConfigDirWin + "mbackup.txt"

// return true if target is a local dir. local dir can be unix style or windows style.
let isLocalTarget (target: string) = target.StartsWith "/" || Regex.IsMatch(target, "^[c-z]:", RegexOptions.IgnoreCase)

// expand user file to mingw64 rsync supported path.
// abc -> /cygdrive/c/Users/<user>/abc
// ^Documents -> expand to Documents path.
// ^Downloads -> expand to Downloads path.
// etc
let expandUserFile (fn: string) =
    let fn =
        let documentsDir =
            Environment.GetFolderPath(Environment.SpecialFolder.MyDocuments)
            |> toMingwPath
            |> ensureDir
        let picturesDir =
            Environment.GetFolderPath(Environment.SpecialFolder.MyPictures)
            |> toMingwPath
            |> ensureDir
        let desktopDir =
            Environment.GetFolderPath(Environment.SpecialFolder.DesktopDirectory)
            |> toMingwPath
            |> ensureDir

        let fn = Regex.Replace(fn, "^My Documents/", documentsDir, RegexOptions.IgnoreCase)
        let fn = Regex.Replace(fn, "^Documents/", documentsDir, RegexOptions.IgnoreCase)
        let fn = Regex.Replace(fn, "^我的文档/", documentsDir)
        let fn = Regex.Replace(fn, "^文档/", documentsDir)
        let fn = Regex.Replace(fn, "^My Pictures/", picturesDir, RegexOptions.IgnoreCase)
        let fn = Regex.Replace(fn, "^Pictures/", picturesDir, RegexOptions.IgnoreCase)
        let fn = Regex.Replace(fn, "^图片/", picturesDir)
        let fn = Regex.Replace(fn, "^Desktop/", desktopDir, RegexOptions.IgnoreCase)
        let fn = Regex.Replace(fn, "^桌面/", desktopDir)
        fn
    if fn.StartsWith("/") then fn
    else userHome + fn

// generate mbackup.list file
let generateMbackupList (logger: Logger) =
    // TODO how to only regenerate if source file have changed? should I bundle GNU make with mbackup?
    // just compare mbackup.list mtime with its source files?
    let mbackupDefaultList = userConfigDirWin + "mbackup-default.list"
    let mbackupLocalList = userConfigDirWin + "local.list"
    let mbackupUserDefaultList = userConfigDirWin + "user-default.list"
    let mbackupList = runtimeDirWin + "mbackup.list"

    // local functions
    let dropEmptyLinesAndComments lines =
        Seq.filter (fun (line: string) -> not (line.TrimStart().StartsWith("#") || line.TrimEnd().Equals(""))) lines
    let readMbackupListFile fn = File.ReadAllLines(fn) |> dropEmptyLinesAndComments

    try
        let defaultListLines = readMbackupListFile mbackupDefaultList |> Seq.map toMingwPath

        let localListLinesMaybe =
            try
                let lines = readMbackupListFile mbackupLocalList |> Seq.map toMingwPath
                (true, lines)
            with
            | :? FileNotFoundException -> (true, Seq.empty)
            | ex ->
                logger.Error "Read mbackupLocalList failed: %s" ex.Message
                (false, Seq.empty)
        match localListLinesMaybe with
        | (false, _) -> failwith "Read mbackup local.list file failed"
        | (true, localListLines) ->
            let userDefaultListLines = readMbackupListFile mbackupUserDefaultList |> Seq.map expandUserFile
            let allLines = Seq.append (Seq.append defaultListLines localListLines) userDefaultListLines
            // For mbackup-default.list and local.list, exclude empty lines and comment lines.
            // skip and give a warning on non-absolute path.
            // For user-default.list, auto prefix user's home dir, auto expand Documents, Downloads etc special folder.
            Directory.CreateDirectory(runtimeDirWin) |> ignore
            File.WriteAllLines(mbackupList, allLines)
            logger.Info
                "mbackup.list file written: %s"
                mbackupList
            true
    with
    | :? System.IO.IOException as ex ->
        logger.Error "Read/write file failed: %s %s" ex.Source ex.Message
        false
    | ex ->
        logger.Error "Read/write mbackup list file failed: %s" ex.Message
        false

exception PrivateKeyNotFoundException of string

let addOptionsForRemoteBackup (results: ParseResults<CLIArguments>) (logger: Logger) (rsyncCmd: string list) =
    let sshExeFile = mbackupProgramDir + "rsync-w64/usr/bin/ssh.exe"
    let sshConfigFile = userHome + ".ssh/config"
    let sshPrivateKeyFile = results.GetResult(Ssh_Key, defaultValue = userHome + ".ssh/id_rsa") |> toMingwPath
    let sshPrivateKeyFileWin = toWinPath sshPrivateKeyFile
    if not (File.Exists(sshPrivateKeyFileWin)) then
        raise (PrivateKeyNotFoundException("ssh private key doesn't exist: " + sshPrivateKeyFileWin))
    else
        let sshConfigFileOption =
            if File.Exists(toWinPath sshConfigFile) then " -F " + sshConfigFile
            else ""

        let rsyncCmd =
            List.append rsyncCmd
                [ sprintf "-e \"'%s'%s -i %s -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null\"" sshExeFile
                      sshConfigFileOption sshPrivateKeyFile ]

        let nodeName = results.GetResult(Node_Name, defaultValue = Net.Dns.GetHostName())
        let remoteLogFile = sprintf "/var/log/mbackup/%s.log" nodeName
        let remoteUser = results.GetResult(Remote_User, defaultValue = Environment.UserName)
        let rsyncCmd = List.append rsyncCmd [ sprintf "--remote-option=--log-file=%s" remoteLogFile ]
        let rsyncCmd = List.append rsyncCmd [ sprintf "--chown=%s:%s" remoteUser remoteUser ]
        rsyncCmd


[<EntryPoint>]
let main argv =
    let errorHandler = ProcessExiter(colorizer = function ErrorCode.HelpText -> None | _ -> Some ConsoleColor.Red)

    let parser = ArgumentParser.Create<CLIArguments>(programName = "mbackup.exe", errorHandler = errorHandler)
    let results = parser.Parse argv

    if results.Contains Version then
        printfn "mbackup %s" versionStr
        Environment.Exit(ExitSuccess)

    let dryRun = results.Contains Dry_Run
    let itemizeChanges = results.Contains Itemize_Changes

    let logger = Logger()

    logger.Info "user config dir: %s" userConfigDirWin
    logger.Info "runtime dir: %s" runtimeDirWin
    logger.Debug "program dir: %s" mbackupProgramDirWin

    let rsyncCmd: string list = []
    let rsyncCmd = appendWhen dryRun rsyncCmd "--dry-run"
    let rsyncCmd = appendWhen itemizeChanges rsyncCmd "-i"
    let rsyncCmd =
        List.append rsyncCmd
            ("-h --stats -togr --delete --delete-excluded --ignore-missing-args".Split [| ' ' |] |> Array.toList)

    let mbackupFile = runtimeDir + "mbackup.list"
    if not (generateMbackupList logger) then failwith "Generate mbackup.list failed"
    let rsyncCmd = List.append rsyncCmd [ sprintf "--files-from=%s" mbackupFile ]

    let excludeFile = userConfigDir + "mbackup-default.exclude"
    let rsyncCmd = List.append rsyncCmd [ sprintf "--exclude-from=%s" excludeFile ]
    let localExcludeFile = userConfigDir + "local.exclude"
    let rsyncCmd = appendWhen (IO.File.Exists localExcludeFile) rsyncCmd (sprintf "--exclude-from=%s" localExcludeFile)

    let localLogFile = runtimeDir + "mbackup.log"
    let rsyncCmd = List.append rsyncCmd [ sprintf "--log-file=%s" localLogFile ]

    // precedence: command line argument > environment variable > config file
    let normalizeTarget target =
        if isLocalTarget target then toMingwPath target
        else target

    let backupTargetMaybe =
        match results.TryGetResult Target with
        | None ->
            let mbackupConfig = WellsConfig(mbackupConfigFile)
            let backupTargetMaybe = mbackupConfig.GetStr("target")
            Option.map normalizeTarget backupTargetMaybe
        | Some backupTarget -> Some(normalizeTarget backupTarget)

    match backupTargetMaybe with
    | None ->
        logger.Error "TARGET is not defined"
        ExitBadParam
    | Some backupTarget ->
        try
            let rsyncCmd =
                if not (isLocalTarget backupTarget) then
                  addOptionsForRemoteBackup results logger rsyncCmd
                else
                  rsyncCmd
            let rsyncCmd = List.append rsyncCmd [ "/" ]
            let rsyncCmd = List.append rsyncCmd [ backupTarget ]
            let rsyncArgs = rsyncCmd |> String.concat " "
            let rsyncExe = mbackupProgramDirWin + "rsync-w64\\usr\\bin\\rsync.exe"
            Directory.CreateDirectory(runtimeDirWin) |> ignore
            Directory.CreateDirectory(userConfigDirWin) |> ignore
            logger.Info
                "Note: if you run the following rsync command yourself, make sure the generated file list (%s) is up-to-date.\n%s"
                mbackupFile (rsyncExe + " " + rsyncArgs)
            let processStartInfo =
                ProcessStartInfo(
                    FileName = rsyncExe,
                    Arguments = rsyncArgs)
            //set HOME dir to prevent ssh.exe can't access /home/<user>/.ssh error.
            try
                processStartInfo.EnvironmentVariables.Add("HOME", userHomeWin)
                setEnv "HOME" userHomeWin
            with
            | :? ArgumentException -> ()    // variable already exists
            | ex -> logger.Warning "set HOME environment variable failed: %A" ex
                // not a critical error, allow program to continue.
            let proc = Process.Start(processStartInfo)
            if proc.WaitForExit Int32.MaxValue then
                logger.Info "mbackup exit"
                proc.ExitCode
            else
                logger.Error "mbackup timed out while waiting for rsync to complete"
                ExitTimeout
        with
        | PrivateKeyNotFoundException msg ->
          logger.Error "%s" msg
          logger.Info
            "backup to remote node requires ssh private key, use --ssh-key <existing_key> option or create ~/.ssh/id_rsa file using ssh-keygen"
          ExitUserError
        | :? IOException as ex ->
            logger.Error "IO Error: %s %s" ex.Source ex.Message
            ExitIOError
        | ex ->
            logger.Error "Unexpected Error: %A" ex
            ExitIOError