#!/usr/bin/env python3

# Copyright 2025 Ping Identity Corporation. All Rights Reserved
#
# This code is to be used exclusively in connection with Ping Identity
# Corporation software or services. Ping Identity Corporation only offers
# such software or services to legal entities who have entered into a
# binding license agreement with Ping Identity Corporation.

# -*- coding: utf-8 -*-

""" Generates user and group entries for applications in the DS pointed
to RCS. For further details, see the README.md file in the same directory.
Do not change the "default" values of the CLI aguments as they are defined
for the baseline use case in pyrock.
"""

# Python imports
import random
import argparse
from ldif3 import LDIFParser

parser = argparse.ArgumentParser()

parser.add_argument(
    "--mode",
    choices=["users", "groups"],
    help="Which mode to run? Default is both users and groups.",
    type=str,
    default=None,
)
parser.add_argument("--num-users", default=100000, help="Number of Users", type=int)
parser.add_argument("--num-apps", default=100, help="Number of Applications", type=int)
parser.add_argument(
    "--input-file",
    default="/results/users.ldif",
    help="Path to input ldif file containing users",
    type=str,
)
parser.add_argument(
    "--output-dir",
    default="./",
    help="Path to output directory for generated ldif files",
    type=str,
)
parser.add_argument(
    "--users-per-app",
    default=1000,
    help="Number of Users per Application",
    type=int,
)
parser.add_argument("--append", action="store_true", help="Append to output file")
parser.add_argument("--random-range", action="store_true", help="Use random range")
parser.add_argument(
    "--groups-per-app",
    default=20,
    help="Number of Groups (entitlments) per Application",
    type=int,
)
parser.add_argument(
    "--min-users-per-group",
    default=10,
    help="Minimum number of Users per Group",
    type=int,
)
parser.add_argument(
    "--max-users-per-group",
    default=1000,  # If --users-per-app is provided then this value is ignored subsequently
    help="Maximum number of Users per Group",
    type=int,
)
# The following arguments are used to generate a specific group for a specific app
# and are not used in the normal case where all groups are generated. This is a hack
# to allow for testing of a specific group without having to generate all groups.
# Do NOT use these arguments in pyrock conf files. They are only for manual testing.
parser.add_argument(
    "--app-number",
    default=-1,
    help="Which application to generate group for? 0 for App0, 1 for App1, etc.",
    type=int,
)
parser.add_argument(
    "--group-name",
    default="groupx-y",
    help="Group name to be used.  For example group0-0, group1-0, etc.",
)
parser.add_argument(
    "--group-size",
    default=100,
    help="Number of members in the group",
    type=int,
)
parser.add_argument(
    "--member-offset",
    default=0,
    help="Start group membership from this offset",
    type=int,
)

args = parser.parse_args()
MODE = args.mode
NUM_USERS = args.num_users  # This value should match "num-entries" in the pyrock conf file
NUM_APPS = args.num_apps  # This value should match number of apps in pyrock conf file
WRITE_FILE_PATH = args.output_dir  # ldif file for DSLdapModifyTask
READ_FILE_PATH = args.input_file  # ldif file from MakeUsersTask
USERS_PER_APP = args.users_per_app  # Number of users per application.
APPEND = True if args.append else False
RANDOM_RANGE = True if args.random_range else False
GROUPS_PER_APP = args.groups_per_app
MIN_USERS_PER_GROUP = args.min_users_per_group
MAX_USERS_PER_GROUP = USERS_PER_APP if USERS_PER_APP else args.max_users_per_group
MANUAL_APP_NUMBER = args.app_number
MANUAL_GROUP_NAME = args.group_name
MANUAL_GROUP_SIZE = args.group_size
MANUAL_MEMBER_OFFSET = args.member_offset
APP0_USERS_PER_GROUP = [  # This is used to determine the number of users in each group for App0
    NUM_USERS,
    NUM_USERS / 2,
    NUM_USERS / 2,
    NUM_USERS / 4,
    NUM_USERS / 10,
]


def main():
    """
    Main function to generate users (ldif) for each mode.  The default behavior is to
    call both modes combined unless explicitly specified in the CLI to run one of them.
    Validation is also performed to ensure output data is consistent with input data.
    """
    match_dn = "ou=people,ou=identities"
    dn_count = 0

    # Validate the users per app argument value to ensure it makes sense
    if USERS_PER_APP:
        min_required = NUM_USERS // NUM_APPS
        if USERS_PER_APP < min_required:
            print(f"ERROR: --users-per-app must be at least {min_required}")
            exit(1)
    # Count the number of users (dn) in the input ldif file
    with open(READ_FILE_PATH, "r") as file:
        for line in file:
            if match_dn in line.lower():
                dn_count += 1

    if dn_count != NUM_USERS:
        print(
            f"ERROR: The number of users in {READ_FILE_PATH} ({dn_count:,}) does not match --num-users ({NUM_USERS:,})."
        )
        exit(1)

    print("=> Number of users:", NUM_USERS)
    print("=> Number of apps:", NUM_APPS)
    print("=> Users per app:", USERS_PER_APP)
    print("=> Calculated apps per user:", NUM_APPS // (NUM_USERS // USERS_PER_APP))
    print("=> Write file path:", WRITE_FILE_PATH)

    if not MODE:
        gen_users_mode()
        gen_groups_mode()

    if MODE == "users":
        gen_users_mode()

    if MODE == "groups":
        gen_groups_mode()


def gen_users_mode():
    """
    Generate user entries for each application.
    """
    app_counter = 1  # Start with App1 as App0 is special
    offset = 0  # Start each application users with this offset
    global mydict  # Dictionary to store first and last names with uid as key
    mydict = {}

    # Build dictionary from ldif file generated on overseer by "MakeUsersTask"
    # This dictionary is subquently used to populate Application users
    # This step is necessary because we want to use the "same" users (as in first and last name)
    # in each app
    populate_dict_from_ldif(READ_FILE_PATH)

    # Write to new ldif file which will be generated on overseer and used to populate
    # the DS instance in the RCS namespace with users for each application using
    # DSLdapModifyTask.  When recon'ed this results in direct assignments in IDM.
    write_file_path = "amIdentityStore-users.ldif" or WRITE_FILE_PATH
    with open(write_file_path, "w") as file:
        while app_counter < NUM_APPS:  # Loop through all Apps
            gen_app_users(file, app_counter, offset)
            offset = offset + USERS_PER_APP
            offset = 0 if offset == NUM_USERS else offset  # recycle users
            app_counter += 1


def gen_groups_mode():
    """
    Generate group entries for each application.
    """
    write_file_path = "amIdentityStore-groups.ldif" or WRITE_FILE_PATH
    app_counter = 0  # Start with App0
    offset = 0  # User offset for each group per App
    mode = "a" if APPEND else "w"
    with open(write_file_path, mode) as file:

        if MANUAL_APP_NUMBER >= 0:  # If a specific app number is provided, generate group for that app and exit
            gen_group_header(file, MANUAL_GROUP_NAME, MANUAL_APP_NUMBER)
            gen_group_members(file, MANUAL_GROUP_SIZE, MANUAL_APP_NUMBER, MANUAL_MEMBER_OFFSET)
            exit(0)

        while app_counter < NUM_APPS:  # Loop through all Apps plus one
            app_group_counter = 0  # Group counter within each App
            group_count = GROUPS_PER_APP

            if app_counter == 0:  # App0
                group_count = len(APP0_USERS_PER_GROUP)  # App0 group count is different

            while app_group_counter < group_count:
                if app_counter == 0:  # App0
                    for range in APP0_USERS_PER_GROUP:
                        grp_name = f"group{app_group_counter}-{app_counter}"
                        gen_group_header(file, grp_name, app_counter)
                        gen_group_members(file, range, app_counter, offset)
                        app_group_counter += 1
                else:
                    grp_name = f"group{app_group_counter}-{app_counter}"
                    if RANDOM_RANGE:
                        range = random.randint(MIN_USERS_PER_GROUP, MAX_USERS_PER_GROUP)
                    else:
                        range = MAX_USERS_PER_GROUP
                    gen_group_header(file, grp_name, app_counter)
                    gen_group_members(file, range, app_counter, offset)

                app_group_counter += 1

            if app_counter != 0:  # if not App0 then increment the offset
                offset = offset + MAX_USERS_PER_GROUP
                offset = 0 if offset == NUM_USERS else offset  # recycle users
            app_counter += 1


def gen_app_users(handle, app_counter, offset):
    """
    Generate user entries for each application.
    These users are used to create direct "assignments" in IDM.
    The RATIO is used to repeat users across applications.
    """

    user_counter = 0  # for each app reset user counter to 0
    while user_counter < USERS_PER_APP:
        uid = f"user.{user_counter+offset}"
        first_name = mydict[uid]["first_name"]
        last_name = mydict[uid]["last_name"]
        mail = mydict[uid]["mail"]  # This is not required by the LDAP schema
        # but is needed for IDM recon co-relation query
        entry = (
            "\n"
            f"dn: uid={uid},ou=app{app_counter},ou=applications,ou=identities\n"
            f"objectclass: top\n"
            f"objectclass: person\n"
            f"objectclass: inetOrgPerson\n"
            f"uid: {uid}\n"
            f"sn: {first_name}\n"
            f"cn: {first_name} {last_name}\n"
            f"mail: {mail}\n"
        )

        handle.write(entry)
        user_counter += 1


def gen_group_header(handle, name, app_counter):

    handle.write(
        "\n"
        + f"dn: ou={name},ou=app{app_counter},ou=applications,ou=identities\n"
        + "objectclass: top\n"
        + "objectClass: groupOfUniqueNames\n"
        + f"ou: {name}\n"
        + f"cn: {name}\n"
    )


def gen_group_members(handle, max, app_counter, offset):
    """
    Generate members (users) for each group.
    handle: file handle for writing
    max: number of users in the group which is determined by "range" in main()
    app_counter: application counter
    offset: user offset for each group per App (except App0).  The offset is used
            to evenly distribute total users across groups within the App.
    """
    user_counter = 0
    while user_counter < max:
        if app_counter == 0:  # App0
            if MANUAL_MEMBER_OFFSET:  # If cli offset is provided
                offset = MANUAL_MEMBER_OFFSET
            else:
                offset = 0
            handle.write(f"uniqueMember: uid=user.{user_counter+offset},ou=people,ou=identities\n")
        else:
            handle.write(
                f"uniqueMember: uid=user.{user_counter+offset},ou=app{app_counter},ou=applications,ou=identities\n"
            )
        user_counter += 1


def populate_dict_from_ldif(path):
    """
    Create dictionary from input ldif file which is generated on the overseer
    """
    sub_suffix = "ou=People,ou=identities"  # case matters
    parser = LDIFParser(open(path, "rb"))

    for dn, entry in parser.parse():
        if dn.endswith(sub_suffix):
            if "givenName" in entry:  # Check if attribute exists in the entry
                givenname = entry["givenName"][0]
            if "sn" in entry:
                sn = entry["sn"][0]
            if "mail" in entry:
                mail = entry["mail"][0]
            if "uid" in entry:
                uid = entry["uid"][0]
                mydict[uid] = {"first_name": givenname, "last_name": sn, "mail": mail}


if __name__ == "__main__":
    main()