#!/usr/bin/python

# Copyright 2022-2023 Univention GmbH
#
# http://www.univention.de/
#
# All rights reserved.
#
# The source code of this program is made available
# under the terms of the GNU Affero General Public License version 3
# (GNU AGPL V3) as published by the Free Software Foundation.
#
# Binary versions of this program provided by Univention to you as
# well as other copyrighted, protected or trademarked materials like
# Logos, graphics, fonts, specific documentations and configurations,
# cryptographic keys etc. are subject to a license agreement between
# you and Univention and not subject to the GNU AGPL V3.
#
# In the case you use this program under the terms of the GNU AGPL V3,
# the program is provided 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 with the Debian GNU/Linux or Univention distribution in file
# /usr/share/common-licenses/AGPL-3; if not, see
# <http://www.gnu.org/licenses/>.

import json
import os
import re
import uuid
from argparse import ArgumentParser

from ldif import LDIFParser, LDIFWriter
from utils import (
    is_school,
    is_school_group,
    is_school_or_school_group,
    is_school_user,
    replace_string_in_ldap_entry,
    status_bar,
)

from univention.config_registry import ConfigRegistry

ucr = ConfigRegistry()
ucr.load()

_NUM_PSEUDONYMS = 30
_PSEUDONYM_ATTRS = [
    "idBrokerPseudonym{}".format(str(i).zfill(4)) for i in range(1, _NUM_PSEUDONYMS + 1)
]
_USER_PASSWORD = "univentionunivention"


class CopySchoolLdif(LDIFParser):
    """
    creates a ldif of a non-exsitent school on the traeger system
    """

    def __init__(self, input, output, old_ou_name, new_ou_name, name_traeger):
        LDIFParser.__init__(self, input)
        self.writer = LDIFWriter(output)
        self.old_ou_name = old_ou_name
        self.new_ou_name = new_ou_name
        self.name_traeger = name_traeger

    def set_user_password(self, entry):  # type: (Dict[str, Any]) -> Dict[str, Any]
        """
        set the password hash of all school users to univentionunivention
        """
        roles = entry.get("ucsschoolRole", [])
        if is_school_user(roles):
            entry["userPassword"] = [
                "{CRYPT}$6$xokxEHm944WkboFx$Hz2lMDsvCNE5Y8JOrFVbYnu4AmX6p0hCQHufDYnkluYkVKudGZ/ac6r4HQvKabsTiVLjuKAaBE2kWmSHMC/ZR0"  # noqa: 501
            ]
            # we don't need krb5Key + ntlm here + it takes time to generate them.
            # entry["sambaNTPassword"] = [ntlm(_user_password)[0]]
            # entry["krb5Key"] = krb5_asn1('{}@{}'.format(entry["uid"][0], REALM), _user_password)

        return entry

    def replace_school_name(self, dn, entry):  # type: (str, Dict[str, Any]) -> Tuple[str, Dict[str, Any]]
        return replace_string_in_ldap_entry(
            dn, entry, self.old_ou_name, self.new_ou_name
        )

    def handle_username(self, dn, entry):  # type: (str, Dict[str, Any]) -> Tuple[str, Dict[str, Any]]
        """
        users receive a reproducable uid in the form of
        ou_name-username
        """
        if entry.get("uid") and "inetOrgPerson" in entry["objectClass"]:
            new_uid = "{}-{}".format(self.new_ou_name, entry["uid"][0])
            assert "test" not in new_uid
            dn = dn.replace(entry["uid"][0], new_uid)
            entry["uid"] = [new_uid]
        return dn, entry

    def handle_group_memberships(self, entry):  # type: (Dict[str, Any]) -> Dict[str, Any]
        """
        replace uids + dns in groups
        """
        if entry.get("memberUid"):
            entry["memberUid"] = [
                "{}-{}".format(self.new_ou_name, member)
                for member in entry["memberUid"]
            ]
        if entry.get("uniqueMember"):
            entry["uniqueMember"] = [
                re.sub(
                    r"uid=([a-zA-Z0-9\-\.\,]+,)",
                    r"uid={}-\1".format(self.new_ou_name),
                    dn,
                )
                for dn in entry["uniqueMember"]
            ]
        return entry

    def handle(self, dn, entry):  # type: (str, Dict[str, Any]) -> None
        # base has to be adjusted with replace_ldap_base
        dn, entry = replace_string_in_ldap_entry(
            dn, entry, ucr.get("ldap/base"), "dc=dummy,dc=base"
        )
        dn, entry = self.replace_school_name(dn, entry)
        dn, entry = self.handle_username(dn, entry)
        entry = self.set_user_password(entry)
        entry["entryUUID"] = [str(uuid.uuid4())]
        entry = self.handle_group_memberships(entry)

        self.writer.unparse(dn, entry)


class CopyIDBrokerLdif(CopySchoolLdif):
    """
    creates a ldif of a non-exsitent school for the id broker system
    input is also the ldif of the trager sytem
    """

    def __init__(self, input, output, old_ou_name, new_ou_name, source_uid):
        CopySchoolLdif.__init__(
            self,
            input=input,
            output=output,
            old_ou_name=old_ou_name,
            new_ou_name=new_ou_name,
            name_traeger=source_uid,
        )
        self.source_uid = source_uid
        self._users = {}
        self._groups = {}
        self._school = {}

    def is_pseudonymized_object(self, entry):  # type: (Dict[str, Any]) -> bool
        roles = entry.get("ucsschoolRole", [])
        return bool(is_school_user(roles) or is_school_or_school_group(roles))

    def handle_transform_to_pseudonymized_object(self, entry):  # type: (Dict[str, Any]) -> Dict[str, Any]
        """
        - set idBrokerPseudonymsType for school-users, -groups and ous
        - set idBrokerOriginIds for school-groups + ous
        - set ucsschoolSourceUID + ucsschoolRecordUID
        - set ucsschoolRecordUID to entryUUID on traeger system
        """
        roles = entry.get("ucsschoolRole", [])
        if self.is_pseudonymized_object(entry):
            entry["objectClass"].append("idBrokerPseudonymsType")
            # setting all pseudonyms - just in case.
            for attr in _PSEUDONYM_ATTRS:
                entry[attr] = [str(uuid.uuid4())]
            if is_school_or_school_group(roles):
                # object class is needed to have the props source + record uid.
                entry["objectClass"].append("idBrokerOriginIds")
            entry["ucsschoolRecordUID"] = entry["entryUUID"]
            entry["ucsschoolSourceUID"] = [self.source_uid]
        return entry

    def _rember_entries(self, entry):  # type: (Dict[str, Any]) -> None
        """
        We save the ids of the test-data in a json to make testing easier.
        """
        roles = entry.get("ucsschoolRole", [])
        if is_school_user(roles):
            uid = entry["uid"][0]
            self._users[uid] = {"id": entry["ucsschoolRecordUID"][0]}
            for key in _PSEUDONYM_ATTRS:
                self._users[uid][key] = entry[key][0]
            self._users[uid]["password"] = _USER_PASSWORD
            self._users[uid]["ucsschoolRole"] = [r.split(":")[0] for r in roles]
        elif is_school_group(roles):
            self._groups[entry["cn"][0]] = {"id": entry["ucsschoolRecordUID"][0]}
            for key in _PSEUDONYM_ATTRS:
                self._groups[entry["cn"][0]][key] = entry[key][0]
            self._groups[entry["cn"][0]]["memberUid"] = entry["memberUid"]
        elif is_school(roles):
            self._school[entry["ou"][0]] = {"id": entry["ucsschoolRecordUID"][0]}
            for key in _PSEUDONYM_ATTRS:
                self._school[entry["ou"][0]][key] = entry[key][0]

    def write_json_files(self, folder, school_name):  # type: (str, str) -> None
        with open(
            os.path.join(folder, "{}_groups.json".format(school_name)), "w"
        ) as fout:
            json.dump(
                {self.source_uid: {school_name: self._groups}},
                fout,
            )
        with open(
            os.path.join(folder, "{}_schools.json".format(school_name)), "w"
        ) as fout:
            json.dump(
                {self.source_uid: self._school},
                fout,
            )
        with open(
            os.path.join(folder, "{}_users.json".format(school_name)), "w"
        ) as fout:
            json.dump(
                {
                    self.source_uid: {
                        school_name: self._users,
                        "provisioning-{}".format(self.source_uid): _USER_PASSWORD,
                    }
                },
                fout,
            )

    def handle(self, dn, entry):  # type: (str, Dict[str, Any]) -> None
        dn, entry = self.replace_school_name(dn, entry)
        entry = self.handle_transform_to_pseudonymized_object(entry)
        self._rember_entries(entry)
        self.writer.unparse(dn, entry)


def create_clone_schools(
    folder, template_ldiff, template_ou_name, traeger_ou_name, name_traeger
):
    """
    creates two ldifs: one for the traeger side + one the id broker side
    + 2 json files.
    """
    school_prefix = "{}-{}".format(name_traeger, traeger_ou_name)
    traeger_output_filename = os.path.join(
        folder, name_traeger, "testdata", "{}.ldif.tmpl".format(school_prefix)
    )
    with (
        open(traeger_output_filename, "wb") as traeger_output,
        open(template_ldiff, "rb") as template_input,
    ):
        parser = CopySchoolLdif(
            input=template_input,
            output=traeger_output,
            old_ou_name=template_ou_name,
            new_ou_name=traeger_ou_name,
            name_traeger=name_traeger,
        )
        parser.parse()

    broker_output_filename = os.path.join(
        folder, "id-broker", "testdata", "{}.ldif.tmpl".format(school_prefix)
    )
    broker_ou_name = "{}-{}".format(name_traeger, traeger_ou_name)
    with (
        open(broker_output_filename, "wb") as broker_output,
        open(traeger_output_filename, "rb") as traeger_input,
    ):
        parser = CopyIDBrokerLdif(
            input=traeger_input,
            output=broker_output,
            old_ou_name=traeger_ou_name,
            new_ou_name=broker_ou_name,
            source_uid=name_traeger,
        )
        parser.parse()
        parser.write_json_files(
            os.path.join(folder, "testdata", "id-broker"), broker_ou_name
        )


def main():
    parser = ArgumentParser()
    parser.add_argument("--num_schools", "-n", required=True)
    parser.add_argument("--template_ldif", "-l", required=True)
    parser.add_argument("--num_traeger", "-t", required=False)
    parser.add_argument("--target_folder", "-o", required=True)
    parser.add_argument("--template_ou_name", "-s", required=False)
    opt = parser.parse_args()

    num_schools = int(opt.num_schools)
    folder = opt.target_folder
    num_traeger = int(opt.num_traeger) or 2
    template_ou_name = opt.template_ou_name or "TEMPLATE"
    template_ldiff = opt.template_ldif

    id_broker_folder = os.path.join(folder, "id-broker", "testdata")
    try:
        os.mkdir(os.path.join(folder, "id-broker"))
        os.mkdir(id_broker_folder)
    except OSError:
        pass
    print("ldifs for id broker will be placed inside: {}".format(id_broker_folder))
    _num_schools_per_traeger = num_schools // num_traeger
    for i in range(1, num_traeger + 1):
        name_traeger = "Traeger{}".format(i)
        print(
            "creating {} schools for {}".format(_num_schools_per_traeger, name_traeger)
        )
        traeger_folder = os.path.join(folder, name_traeger, "testdata")
        try:
            os.mkdir(os.path.join(folder, name_traeger))
            os.mkdir(traeger_folder)
        except OSError:
            pass
        print(
            "ldifs for traeger {} will be placed inside: {}".format(i, traeger_folder)
        )

        for j in range(1, _num_schools_per_traeger + 1):
            traeger_ou_name = "TEST-{}".format(j)
            create_clone_schools(
                folder, template_ldiff, template_ou_name, traeger_ou_name, name_traeger
            )
            status_bar(j, _num_schools_per_traeger)


if __name__ == "__main__":
    main()
