#!/usr/bin/env python
# -*- test-case-name: tests.test_relocate -*-
#
# Utility for relocating Calendar and Contacts data
#
# Copyright (c) 2013 Apple Inc.  All Rights Reserved.
#
# IMPORTANT NOTE:  This file is licensed only for use on Apple-labeled
# computers and is subject to the terms and conditions of the Apple
# Software License Agreement accompanying the package this file is a
# part of.  You may not port this file to another platform without
# Apple's written consent.
from __future__ import print_function

import datetime
import grp
import os
import plistlib
import pwd
import subprocess
import sys
from time import sleep
from twisted.python.filepath import FilePath

LOG = "/var/log/caldavd/relocation.log"
DU = "/usr/bin/du"
DF = "/bin/df"
DITTO = "/usr/bin/ditto"
SUDO = "/usr/bin/sudo"
SERVER_APP_ROOT = "/Applications/Server.app/Contents/ServerRoot"
CALENDARSERVER_CONFIG = "%s/usr/sbin/calendarserver_config" % (SERVER_APP_ROOT,)
CONFIG_ROOT = "/Library/Server/Calendar and Contacts/Config"
COMMAND_FILE = "%s/_relocate_command" % (CONFIG_ROOT,)
PROGRESS_FILE = "%s/_relocate_progress" % (CONFIG_ROOT,)
ERROR_FILE = "%s/_relocate_error" % (CONFIG_ROOT,)
SERVERCTL = "%s/usr/sbin/serverctl" % (SERVER_APP_ROOT,)
LAUNCHCTL = "/bin/launchctl"
PGCTL = "%s/usr/bin/pg_ctl" % (SERVER_APP_ROOT,)



def log(msg):
    try:
        timestamp = datetime.datetime.now().strftime("%b %d %H:%M:%S")
        msg = "relocation: %s %s" % (timestamp, msg)
        with open(LOG, 'a') as output:
            output.write("%s\n" % (msg,)) # so it appears in our log
    except IOError:
        # Could not write to log
        pass



class RelocationError(Exception):
    """
    Data relocation error
    """



def getConfigKeys(keys):
    """
    Ask calendarserver_config for the value of list of keys
    """
    results = {}
    args = [CALENDARSERVER_CONFIG]
    args.extend(keys)

    child = subprocess.Popen(
        args=args,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    # log("Looking up %s" % (keys,))
    output, error = child.communicate()
    # log("Output from calendarserver_config: %s" % (output,))
    if child.returncode:
        log("Error from calendarserver_config: %d, %s" % (child.returncode, error))
    else:
        for line in [l for l in output.split("\n") if "=" in l]:
            key, value = line.strip().split("=", 1)
            results[key] = value

    return results



def writeProgressFile(filename, percentComplete):
    """
    Write a plist with command, percentComplete, and status keys
    @param filename: the full path of the file to write to
    @type filename: C{str}
    @param percentComplete: 0-100
    @type percentComplete: C{int}
    """
    filepath = FilePath(filename)
    d = {
        "command" : "getMoveDataLocationProgress",
        "percentComplete" : percentComplete,
        "status" : "COMPLETED" if percentComplete == 100 else "RUNNING",
    }
    content = plistlib.writePlistToString(d)
    filepath.setContent(content)


def writeErrorFile(filename, msg):
    """
    Write a plist with command and error keys
    @param filename: the full path of the file to write to
    @type filename: C{str}
    @param msg: the error message
    @type msg: C{str}
    """
    filepath = FilePath(filename)
    d = {
        "command" : "getMoveDataLocationProgress",
        "error" : msg.decode("utf-8"),
    }
    content = plistlib.writePlistToString(d)
    filepath.setContent(content)


def readCommandFile(filename):
    """
    Read the given plist file and return the volumePath value
    @param filename: the full path of the file to read from
    @type filename: C{str}
    @return: the value of the volumpePath key within the plist
    @rtype: C{str}
    """
    d = plistlib.readPlist(filename)
    return d.get("volumePath", u"").encode("utf-8")


def removeFile(filename):
    """
    Remove the file or directory (recursive) at filename
    @param filename: the full path of the file or directory to remove
    @type filename: C{str}
    """
    filepath = FilePath(filename)
    if filepath.exists():
        filepath.remove()


def fileExists(filename):
    """
    Return True if file exists, False otherwise
    @param filename: the full path of the file to remove
    @type filename: C{str}
    @return: True if file exists, False otherwise
    @rtype: C{boolean}
    """
    return FilePath(filename).exists()


def copyData(source, dest):
    """
    Ditto source to dest
    """
    args = [DITTO, source, dest]

    child = subprocess.Popen(
        args=args,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    output, error = child.communicate()
    if child.returncode:
        raise(RelocationError("Could not copy from %s to %s: %s" %
            (source, dest, error)))


def setAgentState(enable):
    """
    Use serverctl to enable/disable agent
    """
    args = [SERVERCTL, "enable" if enable else "disable", "service=org.calendarserver.agent"]
    try:
        log("Executing: %s" % (" ".join(args)))
        child = subprocess.Popen(args, stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
        output, error = child.communicate()
        log("Output: %s" % (output,))
        if child.returncode:
            log("Non-zero exit code: %d, %s" % (child.returncode, error))
    except Exception, e:
        log(e)

    if enable:
        # Also launch the agent
        args = [LAUNCHCTL, "start", "org.calendarserver.agent"]
        try:
            log("Executing: %s" % (" ".join(args)))
            subprocess.Popen(args)
        except Exception, e:
            log(e)


def setServiceState(enable):
    """
    Use serverctl to enable/disable calendar/contacts service
    """
    args = [SERVERCTL, "enable" if enable else "disable", "service=org.calendarserver.calendarserver"]
    try:
        log("Executing: %s" % (" ".join(args)))
        child = subprocess.Popen(args, stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
        output, error = child.communicate()
        log("Output: %s" % (output,))
        if child.returncode:
            log("Non-zero exit code: %d, %s" % (child.returncode, error))
    except Exception, e:
        log(e)


def updateConfig(path):
    """
    Update the plist with the new DataRoot
    """
    args = [CALENDARSERVER_CONFIG, "DataRoot=%s" % (path,)]

    child = subprocess.Popen(
        args=args,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    output, error = child.communicate()
    if child.returncode:
        log("Error from calendarserver_config: %d, %s" % (child.returncode, error))


def isPostgresRunning(dataDirectory):
    """
    Ask pg_ctl if postgres is running on the given dataDirectory
    @param dataDirectory: path to data directory
    @type dataDirectory: C{str}
    @return: True if we can connect, False otherwise
    @rtype: C{boolean}
    """
    args = [SUDO, "-u", "calendar", PGCTL, "status", "-D", dataDirectory]

    child = subprocess.Popen(
        args=args,
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    output, error = child.communicate()
    if child.returncode:
        log("Output from pg_ctl: %d, %s" % (child.returncode, error))
        return False
    return True


def exitWithError(msg):
    writeErrorFile(ERROR_FILE, msg)
    removeFile(PROGRESS_FILE)
    sys.exit(0)


def main():

    log("Beginning relocation")

    keys = getConfigKeys(["EnableCalDAV", "EnableCardDAV", "DataRoot",
        "RunRoot", "DatabaseRoot"])
    serviceDataSourcePath = FilePath(keys["DataRoot"])
    log("Source: %s" % (serviceDataSourcePath.path))
    enableCalDAV = keys["EnableCalDAV"] == "True"
    log("EnableCalDAV: %s" % (enableCalDAV,))
    enableCardDAV = keys["EnableCardDAV"] == "True"
    log("EnableCardDAV: %s" % (enableCardDAV,))
    runRoot = keys["RunRoot"]
    log("RunRoot: %s" % (runRoot,))
    databaseRoot = keys["DatabaseRoot"]
    log("DatabaseRoot: %s" % (databaseRoot,))

    writeProgressFile(PROGRESS_FILE, 0)

    # Read and remove command file
    destinationVolume = readCommandFile(COMMAND_FILE)
    log("Command file says destination volume is: %s" % (destinationVolume,))
    removeFile(COMMAND_FILE)

    # Examine destination
    serviceDestination = os.path.join(destinationVolume, "Library", "Server",
        "Calendar and Contacts")
    serviceDestinationPath = FilePath(serviceDestination)
    serviceDataDestinationPath = serviceDestinationPath.child("Data")

    if serviceDataSourcePath.path == serviceDataDestinationPath.path:
        # NO-OP
        writeProgressFile(PROGRESS_FILE, 100)
        log("Destination is the same as source; exiting")
        removeFile(ERROR_FILE)
        return

    if serviceDataDestinationPath.exists():
        msg = "Destination already exists: %s" % (serviceDataDestinationPath.path,)
        log(msg)
        # If destination is boot volume, i.e. destinationVolume == "/", remove
        # any existing data
        if destinationVolume == "/":
            log("Destination is boot volume; removing %s first" %
                (serviceDataDestinationPath.path,))
            serviceDataDestinationPath.remove()
        else:
            log("Destination is not boot volume; exiting")
            exitWithError(msg)

    if not serviceDestinationPath.exists():
        # Calendar and Contacts directory
        log("Creating directory: %s" % (serviceDestination,))
        os.mkdir(serviceDestination)

    try:
        uid = pwd.getpwnam("calendar").pw_uid
        gid = grp.getgrnam("calendar").gr_gid
        os.chown(serviceDestination, uid, gid)
    except:
        pass

    writeProgressFile(PROGRESS_FILE, 25)

    log("Stopping agent")
    setAgentState(False)
    if enableCalDAV or enableCardDAV:
        log("Stopping service")
        setServiceState(False)

    writeProgressFile(PROGRESS_FILE, 50)

    # See if postgres has actually shut down
    dataDirectory = os.path.join(databaseRoot, "cluster.pg")
    tries = 5
    while tries:
        log("Seeing if postgres is still running at %s" % (dataDirectory,))
        if not isPostgresRunning(dataDirectory):
            log("Postgres is not running")
            break
        log("Postgres is still running")
        sleep(5)
        tries -= 1
    else:
        msg = "Postgres never shut down"
        log(msg)
        exitWithError(msg)

    try:
        log("Copying data from %s to %s" %
            (serviceDataSourcePath.path, serviceDataDestinationPath.path))
        copyData(serviceDataSourcePath.path, serviceDataDestinationPath.path)
        log("Done with copy")
    except RelocationError as e:
        exitWithError(str(e))

    writeProgressFile(PROGRESS_FILE, 75)

    log("Updating configuration for DataRoot to %s" %
        (serviceDataDestinationPath.path,))
    updateConfig(serviceDataDestinationPath.path)

    if enableCalDAV or enableCardDAV:
        log("Starting service")
        setServiceState(True)
        sleep(5) # don't start agent at same time as service

    log("Removing error file")
    removeFile(ERROR_FILE)


if __name__ == "__main__":

    if fileExists(PROGRESS_FILE):
        log("Already in progress, exiting")
        sys.exit(0)

    if not fileExists(COMMAND_FILE):
        log("No command file, exiting")
        sys.exit(0)

    try:
        main()
        log("Relocation successful")
    finally:
        log("Removing progress file")
        removeFile(PROGRESS_FILE)
        log("Starting agent")
        setAgentState(True)

