#!/usr/bin/env python
# -*- coding: utf-8 -*-
# BAREOS - Backup Archiving REcovery Open Sourced
#
# Copyright (C) 2015-2015 Planets Communications B.V.
# Copyright (C) 2015-2015 Bareos GmbH & Co. KG
#
# This program is Free Software; you can redistribute it and/or
# modify it under the terms of version three of the GNU Affero General Public
# License as published by the Free Software Foundation, which is
# listed in the file LICENSE.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.
#
# Author: Marco van Wieringen
#
# Bareos python class for LDAP related backup and restore
#

import bareosfd
from bareos_fd_consts import bJobMessageType, bFileType, bRCs, bIOPS
from bareos_fd_consts import bCFs
import BareosFdPluginBaseclass
from StringIO import StringIO
from stat import S_IFREG, S_IFDIR, S_IRWXU
import ldif
import ldap
import ldap.resiter
import ldap.modlist
import time
from calendar import timegm


class BareosFdPluginLDAP(BareosFdPluginBaseclass.BareosFdPluginBaseclass):
    """
    LDAP related backup and restore stuff
    """

    def __init__(self, context, plugindef):
        bareosfd.DebugMessage(
            context,
            100,
            "Constructor called in module %s with plugindef=%s\n"
            % (__name__, plugindef),
        )
        super(BareosFdPluginLDAP, self).__init__(context, plugindef, ["uri", "basedn"])
        self.ldap = BareosLDAPWrapper()

    def parse_plugin_definition(self, context, plugindef):
        """
        Parses the plugin arguments and validates the options.
        """
        bareosfd.DebugMessage(
            context,
            100,
            "parse_plugin_definition() was called in module %s\n" % (__name__),
        )
        super(BareosFdPluginLDAP, self).parse_plugin_definition(context, plugindef)

        return bRCs["bRC_OK"]

    def start_backup_job(self, context):
        """
        Start of Backup Job
        """
        check_option_bRC = self.check_options(context)
        if check_option_bRC != bRCs["bRC_OK"]:
            return check_option_bRC

        return self.ldap.prepare_backup(context, self.options)

    def start_backup_file(self, context, savepkt):
        """
        Defines the file to backup and creates the savepkt.
        """
        bareosfd.DebugMessage(
            context, 100, "BareosFdPluginLDAP:start_backup_file() called\n"
        )

        return self.ldap.get_next_file_to_backup(context, savepkt)

    def start_restore_job(self, context):
        """
        Start of Restore Job
        """
        check_option_bRC = self.check_options(context)
        if check_option_bRC != bRCs["bRC_OK"]:
            return check_option_bRC

        return self.ldap.prepare_restore(context, self.options)

    def start_restore_file(self, context, cmd):
        bareosfd.DebugMessage(
            context,
            100,
            "BareosFdPluginLDAP:start_restore_file() called with %s\n" % (cmd),
        )

        # It can happen that this is the first entry point on a restore
        # and we missed the bEventStartRestoreJob event because we registerd
        # that event after it aready fired. If that is true self.ldap.ld will
        # not be set and we call start_restore_job() from here.
        if not self.ldap.ld:
            return self.start_restore_job(context)
        else:
            return bRCs["bRC_OK"]

    def create_file(self, context, restorepkt):
        """
        Directories are placeholders only we use the data.ldif files
        to get the actual DN of the LDAP record
        """
        bareosfd.DebugMessage(
            context,
            100,
            "BareosFdPluginVMware:create_file() called with %s\n" % (restorepkt),
        )
        if restorepkt.type == bFileType["FT_DIREND"]:
            restorepkt.create_status = bCFs["CF_SKIP"]
        elif restorepkt.type == bFileType["FT_REG"]:
            self.ldap.set_new_dn(restorepkt.ofname)
            restorepkt.create_status = bCFs["CF_EXTRACT"]
        else:
            bareosfd.JobMessage(
                context,
                bJobMessageType["M_FATAL"],
                "Request to restore illegal filetype %s\n" % (restorepkt.type),
            )
            return bRCs["bRC_Error"]

        return bRCs["bRC_OK"]

    def plugin_io(self, context, IOP):
        bareosfd.DebugMessage(
            context,
            100,
            "BareosFdPluginLDAP:plugin_io() called with function %s\n" % (IOP.func),
        )

        if IOP.func == bIOPS["IO_OPEN"]:
            return bRCs["bRC_OK"]

        elif IOP.func == bIOPS["IO_CLOSE"]:
            return bRCs["bRC_OK"]

        elif IOP.func == bIOPS["IO_SEEK"]:
            return bRCs["bRC_OK"]

        elif IOP.func == bIOPS["IO_READ"]:
            if self.ldap.ldif:
                IOP.buf = bytearray(self.ldap.ldif)
                IOP.status = self.ldap.ldif_len
                self.ldap.ldif = None
            else:
                IOP.status = 0
            IOP.io_errno = 0

            return bRCs["bRC_OK"]

        elif IOP.func == bIOPS["IO_WRITE"]:
            self.ldap.ldif = str(IOP.buf)
            self.ldap.ldif_len = IOP.count
            IOP.status = IOP.count
            IOP.io_errno = 0

            return bRCs["bRC_OK"]

    def end_backup_file(self, context):
        bareosfd.DebugMessage(
            context, 100, "BareosFdPluginLDAP:end_backup_file() called\n"
        )

        return self.ldap.has_next_file(context)

    def end_restore_file(self, context):
        bareosfd.DebugMessage(
            context, 100, "BareosFdPluginLDAP:end_restore_file() called\n"
        )

        return self.ldap.restore_entry(context)


# Overload ldap class and ldap.resiter class for bulk retrieval
class BulkLDAP(ldap.ldapobject.LDAPObject, ldap.resiter.ResultProcessor):
    pass


class BareosLDAPWrapper:
    """
    LDAP specific class.
    """

    def __init__(self):
        self.ld = None
        self.resultset = None
        self.ldap_entries = None
        self.file_to_backup = None
        self.dn = None
        self.entry = None
        self.ldif = None
        self.ldif_len = None
        self.unix_create_time = None
        self.unix_modify_time = None
        self.msg_id = None

    def connect_and_bind(self, context, options, bulk=False):
        """
        Bind to LDAP URI using the given authentication tokens
        """
        if bulk:
            self.ld = BulkLDAP(options["uri"])
        else:
            self.ld = ldap.initialize(options["uri"])

        try:
            self.ld.protocol_version = ldap.VERSION3
            if "bind_dn" in options and "password" in options:
                self.ld.simple_bind_s(options["bind_dn"], options["password"])
            else:
                self.ld.simple_bind_s("", "")
        except ldap.INVALID_CREDENTIALS:
            bareosfd.JobMessage(
                context,
                bJobMessageType["M_FATAL"],
                "Failed to bind to LDAP uri %s\n" % (options["uri"]),
            )

            return bRCs["bRC_Error"]
        except ldap.LDAPError as e:
            if type(e.message) == dict and "desc" in e.message:
                bareosfd.JobMessage(
                    context,
                    bJobMessageType["M_FATAL"],
                    "Failed to bind to LDAP uri %s: %s\n"
                    % (options["uri"], e.message["desc"]),
                )
            else:
                bareosfd.JobMessage(
                    context,
                    bJobMessageType["M_FATAL"],
                    "Failed to bind to LDAP uri %s: %s\n" % (options["uri"], e),
                )

            return bRCs["bRC_Error"]

        bareosfd.DebugMessage(context, 100, "connected to LDAP server\n")

        return bRCs["bRC_OK"]

    def prepare_backup(self, context, options):
        """
        Prepare a LDAP backup
        """
        connect_bRC = self.connect_and_bind(context, options, True)
        if connect_bRC != bRCs["bRC_OK"]:
            return connect_bRC

        bareosfd.DebugMessage(
            context,
            100,
            "Creating search filter and attribute filter to perform LDAP search\n",
        )

        # Get a subtree
        searchScope = ldap.SCOPE_SUBTREE

        # See if there is a specific search filter otherwise use the all object filter.
        if "search_filter" in options:
            searchFilter = options["search_filter"]
        else:
            searchFilter = r"(objectclass=*)"

        # Get all user attributes and createTimestamp + modifyTimestamp
        attributeFilter = ["*", "createTimestamp", "modifyTimestamp"]

        try:
            # Asynchronous search method
            self.msg_id = self.ld.search(
                options["basedn"], searchScope, searchFilter, attributeFilter
            )
        except ldap.LDAPError as e:
            if type(e.message) == dict and "desc" in e.message:
                bareosfd.JobMessage(
                    context,
                    bJobMessageType["M_FATAL"],
                    "Failed to execute LDAP search on LDAP uri %s: %s\n"
                    % (options["uri"], e.message["desc"]),
                )
            else:
                bareosfd.JobMessage(
                    context,
                    bJobMessageType["M_FATAL"],
                    "Failed to execute LDAP search on LDAP uri %s: %s\n"
                    % (options["uri"], e),
                )

            return bRCs["bRC_Error"]

        bareosfd.DebugMessage(context, 100, "Successfully performed LDAP search\n")

        return bRCs["bRC_OK"]

    def prepare_restore(self, context, options):
        """
        Prepare a LDAP restore
        """
        connect_bRC = self.connect_and_bind(context, options)
        if connect_bRC != bRCs["bRC_OK"]:
            return connect_bRC

        return bRCs["bRC_OK"]

    def to_unix_timestamp(self, context, timestamp):
        if timestamp:
            unix_time = timegm(
                time.strptime(timestamp.replace("Z", "GMT"), "%Y%m%d%H%M%S%Z")
            )
            return int(unix_time)
        else:
            return None

    def get_next_file_to_backup(self, context, savepkt):
        """
        Find out the next file that should be backed up
        """
        # When file_to_backup is not None we should return the LDIF.
        if self.file_to_backup:
            # Remove some attributes from entry before creating the LDIF.
            ignore_attribute = ["createTimestamp", "modifyTimestamp"]

            keys = self.entry.keys()
            for value in keys:
                if value in ignore_attribute:
                    del self.entry[value]

            # Dump the content of the LDAP entry as LDIF text
            ldif_dump = StringIO()
            ldif_out = ldif.LDIFWriter(ldif_dump)
            ldif_out.unparse(self.dn, self.entry)
            self.ldif = ldif_dump.getvalue()
            self.ldif_len = len(self.ldif)
            ldif_dump.close()

            statp = bareosfd.StatPacket()
            statp.mode = S_IRWXU | S_IFREG
            statp.size = self.ldif_len
            if self.unix_create_time:
                statp.ctime = self.unix_create_time
            if self.unix_modify_time:
                statp.mtime = self.unix_modify_time

            savepkt.statp = statp
            savepkt.type = bFileType["FT_REG"]
            savepkt.fname = self.file_to_backup + "/data.ldif"
            # Read the content of a file
            savepkt.no_read = False

            # On next run we need to get next entry from result set.
            self.file_to_backup = None
        else:
            # If we have no result set get what the LDAP search returned as resultset.
            if self.resultset is None:
                self.resultset = self.ld.allresults(self.msg_id)
                # Try to get the first result set from the query,
                # if there is nothing return an error.
                try:
                    res_type, res_data, res_msgid, res_controls = self.resultset.next()
                    self.ldap_entries = res_data
                except ldap.NO_SUCH_OBJECT:
                    return bRCs["bRC_Error"]
                except StopIteration:
                    return bRCs["bRC_Error"]

            # Get the next entry from the result set.
            if self.ldap_entries:
                self.dn, self.entry = self.ldap_entries.pop(0)

                if self.dn:
                    # Extract the createTimestamp and modifyTimestamp and
                    # convert it to an UNIX timestamp
                    self.unix_create_time = None
                    try:
                        createTimestamp = self.entry["createTimestamp"][0]
                    except KeyError:
                        pass
                    else:
                        self.unix_create_time = self.to_unix_timestamp(
                            context, createTimestamp
                        )

                    self.unix_modify_time = None
                    try:
                        modifyTimestamp = self.entry["modifyTimestamp"][0]
                    except KeyError:
                        pass
                    else:
                        self.unix_modify_time = self.to_unix_timestamp(
                            context, modifyTimestamp
                        )

                    # Convert the DN into a PATH e.g. reverse the elements.
                    dn_sliced = self.dn.split(",")
                    self.file_to_backup = "@LDAP" + "".join(
                        ["/" + element for element in reversed(dn_sliced)]
                    )

                    statp = bareosfd.StatPacket()
                    statp.mode = S_IRWXU | S_IFDIR
                    if self.unix_create_time:
                        statp.ctime = self.unix_create_time
                    if self.unix_modify_time:
                        statp.mtime = self.unix_modify_time

                    savepkt.statp = statp
                    savepkt.type = bFileType["FT_DIREND"]
                    savepkt.fname = self.file_to_backup
                    # A directory has a link field which contains the fname + a trailing '/'
                    savepkt.link = self.file_to_backup + "/"
                    # Don't read the content of a directory
                    savepkt.no_read = True

        return bRCs["bRC_OK"]

    def set_new_dn(self, fname):
        path = fname[: -len("/ldif.data")]
        # Convert the PATH into a DN e.g. reverse the elements.
        path_sliced = path.split("/")
        new_dn = "".join([element + "," for element in reversed(path_sliced)])
        # Remove the ',@LDAP,' in the DN which is an internal placeholder
        self.dn = new_dn.replace(",@LDAP,", "")

    def has_next_file(self, context):
        # See if we are currently handling the LDIF file or
        # if there is still data in the result set
        if self.file_to_backup or self.ldap_entries:
            bareosfd.DebugMessage(context, 100, "has_next_file(): returning bRC_More\n")
            return bRCs["bRC_More"]
        else:
            # See if there are more result sets.
            if self.resultset:
                try:
                    # Get the next result set
                    res_type, res_data, res_msgid, res_controls = self.resultset.next()
                    self.ldap_entries = res_data

                    # We expect something to be in the result set but better check
                    if self.ldap_entries:
                        bareosfd.DebugMessage(
                            context, 100, "has_next_file(): returning bRC_More\n"
                        )
                        return bRCs["bRC_More"]
                except StopIteration:
                    pass

        return bRCs["bRC_OK"]

    def restore_entry(self, context):
        # Restore the entry
        if self.ldif:
            # Parse the LDIF back to an attribute list
            ldif_dump = StringIO(str(self.ldif))
            ldif_parser = ldif.LDIFRecordList(ldif_dump, max_entries=1)
            ldif_parser.parse()
            dn, entry = ldif_parser.all_records[0]
            ldif_dump.close()

            if self.dn != dn:
                bareosfd.JobMessage(
                    context,
                    bJobMessageType["M_INFO"],
                    "Restoring original DN %s as %s\n" % (dn, self.dn),
                )

            if dn:
                if self.ld:
                    # Try adding the entry
                    add_ldif = ldap.modlist.addModlist(entry)
                    try:
                        self.ld.add_s(self.dn, add_ldif)
                    except ldap.LDAPError as e:
                        # Delete the original DN
                        try:
                            self.ld.delete_s(self.dn)
                            self.ld.add_s(self.dn, add_ldif)
                        except ldap.LDAPError as e:
                            if type(e.message) == dict and "desc" in e.message:
                                bareosfd.JobMessage(
                                    context,
                                    bJobMessageType["M_ERROR"],
                                    "Failed to restore LDAP DN %s: %s\n"
                                    % (self.dn, e.message["desc"]),
                                )
                            else:
                                bareosfd.JobMessage(
                                    context,
                                    bJobMessageType["M_ERROR"],
                                    "Failed to restore LDAP DN %s: %s\n" % (self.dn, e),
                                )
                            self.ldif = None
                            return bRCs["bRC_Error"]
                else:
                    bareosfd.JobMessage(
                        context,
                        bJobMessageType["M_ERROR"],
                        "Failed to restore LDAP DN %s no writable binding to LDAP exists\n"
                        % (self.dn),
                    )
                    self.ldif = None
                    return bRCs["bRC_Error"]

            # Processed ldif
            self.ldif = None

        return bRCs["bRC_OK"]

    def cleanup(self, context):
        if self.ld:
            self.ld.unbind_s()

        return bRCs["bRC_OK"]


# vim: ts=4 tabstop=4 expandtab shiftwidth=4 softtabstop=4
