#!/usr/bin/python
# Chopu AIMA upgrade script

import isi.config.auth as authconf
import isi.config.lwreg as lwreg
import isi.config.mcp as mcp
import isi.config.mxd as mxd
import isi.config.xml as xml
import isi.gconfig as gcfg
import isi.upgrade.authcheck as authcheck
import isi.upgrade.authusers as authusers

import os
import sys
import time
import logging
import syslog
import isi.auth as auth
from isi.app.lib import safelockfile
from isi.ui.lib import nsswitch
import isi.sys.version as sysver
import errno

IFS_SAMBA_PATH = "/ifs/.ifsvar/samba"
IFS_BACKUP_PATH = "/ifs/.ifsvar/tmp"
IFS_LIKEWISE_DB_PATH = "/ifs/.ifsvar/likewise/db"
ETC_PATH = "/etc"

OVERRIDE_LSASSD_XML = "/etc/mcp/override/lsassd.xml"

OVERRIDE_AUTH_XML = "/etc/mcp/override/auth"

LOG_FILE = os.path.join(IFS_BACKUP_PATH, "chopu-auth-upgrade.log")
AUTH_UPGRADE_FILE = os.path.join(IFS_BACKUP_PATH, "mavericks-auth-upgraded")

OVERRIDE_NIS_XML = "/etc/mcp/override/nis.xml"
NIS_XML = "/etc/mcp/sys/nis.xml"

# 5.5.x and 6.0.x machine password location
PSTORE_PATH = os.path.join(IFS_LIKEWISE_DB_PATH, "pstore.db")

# 5.5.x and 6.0.x local users/groups and group membership location
LSASS_LOCAL_PATH = os.path.join(IFS_LIKEWISE_DB_PATH, "lsass-local.db")

# 5.5.x and 6.0.x idmap location
LSASS_IDMAP_PATH = os.path.join(IFS_LIKEWISE_DB_PATH, "idmap.db")

MAX_GCONF_COMMIT_RETRIES = 5

logger = logging.getLogger("chopu_auth_upgrade")
authusers.set_logger(logger)

######################################################################
#  Upgrade Steps
##########################################################################
def _needs_config_upgrade(aconf, version):
    ''' Check the config and return True if upgrade is necessary. '''
    return aconf.compare_config_version(version) < 0

def step_upgrade_config650():
    # bug 111806: turn off a gconfig validation error during this upgrade step
    os.environ['GC_NO_NAME_MISMATCH_ERRORS'] = '1'
    retries = 0
    while True:
        cfg = gcfg.open(gcfg.gcfg_main)
        ctx = cfg.ctx_new()
        # first check latest version for config.
        v2conf = authconf.GlobalConfig(context=ctx)
        path = v2conf.get_config_version_path()
        try:
            ctx.read(path)
        except gcfg.VersionError:
            # not at version 2
            pass
        else:
            if not _needs_config_upgrade(v2conf, lwreg.V6_5_0):
                logger.debug("Config already upgraded to version %s" %
                             v2conf.get_config_version())
                break

        ctx.read_create("registry:1")
        aconf = authcheck.Auth65Config(context=ctx)
            
            
        # All nodes will attempt the upgrade.  As soon as one
        # finishes, give up.
        if not _needs_config_upgrade(aconf, lwreg.V6_5_0):
            logger.debug("Config already upgraded to version %s" %
                               aconf.get_config_version())
            break
        logger.debug("upgrading from version %s to version 6.5.0" %
                    aconf.get_config_version())
        # Attempt to update 
        try:
            authcheck.parse_lsassd_xml(aconf._context,
                lsassd_path=OVERRIDE_LSASSD_XML, logger=logger)
        except IOError as e:
            if e.errno == errno.ENOENT:
                pass
            else:
                del os.environ['GC_NO_NAME_MISMATCH_ERRORS']
                raise

        upgraded = sysver.Node_Status().upgraded_across_chopu()
        # AIMA is Opt-in for upgrades.
        if upgraded:
            aconf.set_single_param("on-disk-identity", "unix")
        
        # Update the version to prevent other nodes from trying.
        aconf.set_config_version(lwreg.V6_5_0)
        logger.info("Config updated to version %s" % lwreg.V6_5_0)

        try:
            aconf.commit()
            logger.info("Upgrade commit to %s successful." % lwreg.V6_5_0)
        except gcfg.RecoverableCommitError, e:
            if retries >= MAX_GCONF_COMMIT_RETRIES:
                logger.info("Upgrade config commit failed: %s. " % str(e))
                del os.environ['GC_NO_NAME_MISMATCH_ERRORS']
                raise
            logger.info("Upgrade config commit failed: %s. "
                              "Retrying upgrade." % str(e))
            retries += 1
            continue
        except gcfg.CommitError, e:
            logger.info("Upgrade config commit failed: %s. " % str(e))
            del os.environ['GC_NO_NAME_MISMATCH_ERRORS']
            raise
        break

    del os.environ['GC_NO_NAME_MISMATCH_ERRORS']
    return True


def step_upgrade_config651():
    os.environ['GC_NO_NAME_MISMATCH_ERRORS'] = '1'
    retries = 0
    while True:
        cfg = gcfg.open(gcfg.gcfg_main)
        ctx = cfg.ctx_new()
        # first check latest version for config.
        v2conf = authconf.GlobalConfig(context=ctx)
        path = v2conf.get_config_version_path()
        try:
            ctx.read(path)
        except gcfg.VersionError:
            # not at version 2
            pass
        else:
            if not _needs_config_upgrade(v2conf, lwreg.V6_5_1):
                logger.debug("Config already upgraded to version %s" %
                             v2conf.get_config_version())
                break

        ctx.read_create("registry:1")
        aconf = authcheck.Auth65Config(context=ctx)

        # All nodes will attempt the upgrade.  As soon as one
        # finishes, give up.
        if not _needs_config_upgrade(aconf, lwreg.V6_5_1):
            logger.debug("Config already upgraded to version %s" %
                               aconf.get_config_version())
            break

        logger.debug("upgrading from version %s to version 6.5.1" %
                    aconf.get_config_version())

        # Load existing settings
        try:
            nisdoc = mxd.loadmxd(NIS_XML, OVERRIDE_NIS_XML)
        except mxd.error, e:
            raise xml.XMLConfigError(e)

        nisroot = nisdoc.getroot()

        # Domain name:
        domain = nisroot.getchildnamed('domain')
        if domain is not None:
            domainName = domain.gettext()
        else:
            domainName = ''

        # Servers:
        servers = []
        nisroot.reset()
        try:
            servers_ref = nisroot.getchildnamed('servers')
            servers = map(lambda a: a.gettext(), servers_ref.getchildrennamed('server'))
        except AttributeError:
            pass

        # NSSwitch
        nss = nsswitch.NSSwitch()

        # Only disable NIS if there are servers, a domain configured,
        # and NIS was disabled.  Otherwise, ensure the default
        # of enabled.
        if servers and domainName and not nss.sourceExists('nis'):
            aconf.set_single_param("nis-enabled", "false")
        else:
            aconf.set_single_param("nis-enabled", "true")

        # Update auth settings
        if servers:
            aconf.set_single_param("nis-server", ','.join(servers))
        if domainName:
            aconf.set_single_param("nis-domain", domainName)

        # Update the version to prevent other nodes from trying.
        aconf.set_config_version(lwreg.V6_5_1)
        logger.info("Config updated to version %s" % lwreg.V6_5_1)

        try:
            aconf.commit()
            logger.info("Upgrade commit to %s successful." % lwreg.V6_5_1)
        except gcfg.RecoverableCommitError, e:
            if retries >= MAX_GCONF_COMMIT_RETRIES:
                logger.info("Upgrade config commit failed: %s. " % str(e))
                del os.environ['GC_NO_NAME_MISMATCH_ERRORS']
                raise
            logger.info("Upgrade config commit failed: %s. "
                              "Retrying upgrade." % str(e))
            retries += 1
            continue
        except gcfg.CommitError, e:
            logger.info("Upgrade config commit failed: %s. " % str(e))
            del os.environ['GC_NO_NAME_MISMATCH_ERRORS']
            raise

        break

    del os.environ['GC_NO_NAME_MISMATCH_ERRORS']
    return True

def step_upgrade_disableNIS():
    os.environ['GC_NO_NAME_MISMATCH_ERRORS'] = '1'
    cfg = gcfg.open(gcfg.gcfg_main)
    ctx = cfg.ctx_new()
    try:
        ctx.read("registry:1")
    except gcfg.VersionError:
        # not at version 1
        del os.environ['GC_NO_NAME_MISMATCH_ERRORS']
        return True

    aconf = authcheck.Auth65Config(context=ctx)

    # If we've upgraded to 6.5.1 auth settings, disable NIS (if necessary)
    if _needs_config_upgrade(aconf, lwreg.V6_5_1):
        logger.debug("Config at version %s.  Needs to be 6.5.1" %
                               aconf.get_config_version())
        del os.environ['GC_NO_NAME_MISMATCH_ERRORS']
        return False

    nss = nsswitch.NSSwitch()
    mcpconfig = mcp.mcpconfig()

    if nss.sourceExists('nis'):
        logger.info("Disabling NIS source.")
        nss.deleteSource('nis')
        logger.info("Adding LSASS to netgroup.")
        nss.addSourceToDB('lsass', 'netgroup')
        nss.save()

    if mcpconfig.get('nis') and mcpconfig.get('nis')['enabled']:
        logger.info("Disabling NIS MCP.")
        mcpconfig.disable('nis')

    del os.environ['GC_NO_NAME_MISMATCH_ERRORS']
    return True

def step_start_auth():
    try:
        auth.ping()
        return True
    except Exception:
        pass

    # we need to start up auth
    mcpconfig = mcp.mcpconfig()
    mcpconfig.showall()
    mcpconfig.monitor('auth')
    mcpconfig.enable('auth')
    # give some time for mcp to react
    mcpconfig.runAction('lwsm', 'start-lwsmd')
    time.sleep(3)
    mcpconfig.runAction('auth', 'start-lsass')

    # Wait up to 2 minutes for auth to start.  If it is still not responsive at
    # the end of the time delay, we continue anyway as to not hold up the boot
    # process.
    start_time = time.time()
    seconds = 0
    wait = 1
    retried = 0
    while seconds <= 120:
        try:
            auth.ping()
            break
        except auth.ConnectionRefusedError:
            if seconds >= 10 and retried < 1:
                # give lsass startup another shot
                retried = 1
                logger.warning("retrying lsass")
                mcpconfig.runAction('auth', 'start-lsass')
                
            if seconds >= 60 and retried < 2:
                retried = 2
                # If we still can't ping lsassd after one minute restart
                # the service
                logger.warning("retrying lwsm/lsass")
                mcpconfig.runAction('lwsm', 'start-lwsmd')
                time.sleep(30)
                mcpconfig.runAction('auth', 'start-lsass')

        time.sleep(wait)
        seconds = time.time() - start_time
    else:
        # Runs when timeout expires
        logger.warning("Auth daemons were unable to initialize or contact "
                       "the primary domain controller, continuing with upgrade "
                       "anyway.")

    return True

def try_step_rename(prev_step, cur_step, func, base_path):
    '''Helper function to handle renaming of files
        as we progress through substeps'''
    prev_path = "%s%s" % (base_path, prev_step and ("-%s" % prev_step) or "")
    passed_path = "%s-%s" % (base_path, cur_step)
    failed_path = "%s-failed" % passed_path
    if os.path.isfile(prev_path):
        logger.info("Starting substep %s" % cur_step)
        try:
            func(prev_path)
        except:
            logger.warning("Substep %s failed.  "
                "Renaming file to %s" % (cur_step, failed_path))
            os.rename(prev_path, failed_path)
            raise
        else:
            os.rename(prev_path, passed_path)
    elif os.path.isfile(failed_path):
        logger.info("Re-starting substep %s" % cur_step)
        try:
            func(failed_path)
        except:
            logger.warning("Substep %s failed." % cur_step)
            raise
        else:
            os.rename(failed_path, passed_path)        
    else:
        logger.info("Skipping already completed substep %s" % cur_step)


def step_upgrade_pstore(pstore_path=PSTORE_PATH):
    ''' Import machine password if cluster is joined into a domain. '''
    logger.info("Start reading machine password from pstore")

    try:
        try_step_rename(None, "import_machine_password",
            authusers.import_machpass_from_pstore, pstore_path)

        try_step_rename("import_machine_password", "finished",
            lambda path: None, pstore_path)
    except Exception, e:
        logger.warning("Failed to read machine password from pstore (%s).  "
            "Error was: %s" % (pstore_path, e))
        success = False
    else:
        if os.path.isfile("%s-finished" % pstore_path):
            logger.info("Finished reading machine password from pstore.  "
                        "Database saved as %s-finished." % pstore_path)
        else:
            logger.info("Finished reading machine password from pstore.")
        success = True

    return success


def step_upgrade_idmap(idmap_path=LSASS_IDMAP_PATH,
        lsassd_xml=OVERRIDE_LSASSD_XML):
    '''Get the idmap range, machine sid, and mappings from the idmap.db'''
    logger.info("Start importing lsass idmap database from %s" %
        idmap_path)

    try:
        try_step_rename(None, "set_machine_sid",
            authusers.import_machine_sid, idmap_path)

        try_step_rename("set_machine_sid", "import_idmap",
            lambda idmap_db: authusers.import_idmap_db(idmap_db, lsassd_xml),
                idmap_path)

        try_step_rename("import_idmap", "finished",
            lambda path: None, idmap_path)
    except Exception, e:
        logger.warning("Failed to import lsass idmap database (%s).  "
            "Error was: %s" % (idmap_path, e))
        success = False
    else:
        if os.path.isfile("%s-finished" % idmap_path):        
            logger.info("Finished importing lsass idmap database.  "
                        "Database saved as %s-finished" % idmap_path)
        else:
            logger.info("Finished importing lsass idmap database.")            
        success = True

    return success


def step_upgrade_lsass_local(local_path=LSASS_LOCAL_PATH):
    '''Get all the users, groups, and membership from lsass-local.db'''
    logger.info("Start importing users and groups from %s" %
        local_path)

    try:
        try_step_rename(None, "import_users_and_groups",
            authusers.import_from_lsass_local, local_path)

        try_step_rename("import_users_and_groups", "import_group_membership",
            authusers.import_group_membership_from_lsass_local,
            local_path)

        try_step_rename("import_group_membership", "finished",
            lambda path: None, local_path)
    except Exception, e:
        logger.warning(
            "Failed to import local users and groups database (%s).  "
            "Error was: %s" % (local_path, e))
        success = False
    else:
        if os.path.isfile("%s-finished" % local_path):
            logger.info("Finished importing users, groups and membership.  "
                        "Database saved as %s-finished" % local_path)
        else:
            logger.info("Finished importing users, groups and membership.")            
        success = True

    return success


upgrade_steps = [
    [
        # step 0 is empty
    ],
    [
        ("upgrade_config650", step_upgrade_config650),
        ("upgrade_config651", step_upgrade_config651),
        ("upgrade_disableNIS", step_upgrade_disableNIS),
        ("upgrade_idmap", step_upgrade_idmap),
    ],
    [
        ("start_auth", step_start_auth),
        ("upgrade_users_and_groups", step_upgrade_lsass_local),
        ("import_machine_passwd", step_upgrade_pstore),
    ]
]

def main(argv):
    upgrade_pass = 1
    if (len(argv) > 1):
        upgrade_pass = int(argv[1])
        
    try:
        logger.addHandler(logging.FileHandler(LOG_FILE, 'a'))
    except EnvironmentError, e:
        syslog.syslog(syslog.LOG_NOTICE,
                      "Error running auth upgrade script, %s"%e)
        return 0
    logger.setLevel(logging.DEBUG)

    try:
        os.chmod(LOG_FILE, 0640)
    except OSError, e:
        # log chmod errors and continue
        syslog.syslog(syslog.LOG_NOTICE,
                      "Error chmod'ing %s (errno: %d)" % (LOG_FILE, e.errno))

    nstat = sysver.Node_Status()
    prevNodeVersion = nstat.get_previous_version()

    if upgrade_pass <= 0 or upgrade_pass >= len(upgrade_steps):
        logger.warning("Bad Chopu Auth Upgrade step (%d)" % upgrade_pass)
        return 1

    # if we are reimaging a new node, skip Chopu upgrade
    if prevNodeVersion is None:
        logger.info("%s: Skipped Chopu Auth Upgrade step %d: reimage of new node" %
                    (time.ctime(), upgrade_pass))
        return 0

    # if we are updating from a 7.0 (Mavericks) node, skip Chopu upgrade
    if prevNodeVersion >= 0x606000000000000L:
        logger.info("%s: Skipped Chopu Auth Upgrade step %d: node is version >= 7.0" %
                    (time.ctime(), upgrade_pass))
        return 0

    # skip starting up lsass if we can avoid it. 
    if upgrade_pass == 2 and prevNodeVersion >= 0x605000000000000L:
        logger.info("%s: Skipped Chopu Auth Upgrade step %d: "
                    "node is version >= 6.5" % (time.ctime(), upgrade_pass))
        return 0

    # We need to block all other nodes from continue until we've finished
    # the upgrade
    try:
        sentry = safelockfile(AUTH_UPGRADE_FILE, 'w')
    except EnvironmentError, e:
        syslog.syslog(syslog.LOG_NOTICE,
                      "Error running auth upgrade script, %s"%e)
        return 0


    if (upgrade_pass == 1):
        logger.info("\n\n")
    logger.info("%s: Chopu Auth Upgrade Step %d" %
                (time.ctime(), upgrade_pass))
    if (upgrade_pass == 1):
        logger.info("%s: Node Status: upgrade version from '%s' to '%s'" %
                    (time.ctime(),
                     nstat.get_previous_version_string(),
                     nstat.get_current_version_string()))


    for step in upgrade_steps[upgrade_pass]:

        logger.info("%s: STARTED step %s"%(time.ctime(), step[0]))
        try:
            result = step[1]()
            if not result:
                logger.error("%s: FAILED step %s" %(time.ctime(), step[0]))
            else:
                logger.info("%s: SUCCEEDED step %s" %(time.ctime(), step[0]))
        except Exception, e:
            import traceback
            logger.error("%s: EXCEPTION on step %s: %s" %(time.ctime(),
                step[0], traceback.format_exc()))
            result = False

    sentry.close()

    return 0

if __name__ == "__main__":
    sys.exit(main(sys.argv))
