import os
import random
import string
import pwd
import grp

# Function to create a random folder name
def random_name(length=6):
    return ''.join(random.choices(string.ascii_lowercase, k=length))

# Function to create a random file name with extension
def random_file_name(extension='txt'):
    return random_name() + f'.{extension}'

# Function to assign random permissions for files
def random_chmod():
    return oct(random.choice([
        0o644,  # rw-r--r--
        0o755,  # rwxr-xr-x
        0o600,  # rw-------
        0o700,  # rwx------
        0o444,  # r--r--r--
        0o666,  # rw-rw-rw-
        0o777,  # rwxrwxrwx
        0o555,  # r-xr-xr-x
        0o222,  # --w--w--w-
        0o111   # --x--x--x
    ]))

# Function to set ownership based on the user (not root yet)
def set_ownership(path, uid, gid):
    os.chown(path, uid, gid)

# Function to randomly assign ownership to root after creation
def randomly_assign_root_ownership(base_dir, uid, gid, root_uid, root_gid):
    for root, dirs, files in os.walk(base_dir):
        for file_name in files:
            file_path = os.path.join(root, file_name)
            if random.choice([True, False]):
                # Assign ownership to root
                os.chown(file_path, root_uid, root_gid)
                print(f'Set ownership to root for {file_path}')

# Function to create the folder structure
def create_random_folder_structure(base_dir, max_depth=3, current_depth=0, uid=None, gid=None):
    if current_depth < max_depth:
        num_folders = random.randint(1, 3)
        for _ in range(num_folders):
            folder_name = random_name()
            folder_path = os.path.join(base_dir, folder_name)
            os.makedirs(folder_path, exist_ok=True)
            print(f'Created directory: {folder_path}')

            # Optionally create some files in this folder
            num_files = random.randint(1, 3)
            for _ in range(num_files):
                file_name = random_file_name()
                file_path = os.path.join(folder_path, file_name)
                with open(file_path, 'w') as f:
                    f.write('This is a dummy file.\n')
                print(f'Created file: {file_path}')

                # Set permissions and ownership for the file
                os.chmod(file_path, int(random_chmod(), 8))  # Random permissions for files
                set_ownership(file_path, uid, gid)  # Set ownership for files to user
                print(f'Set ownership and permissions for file: {file_path}')

            # Recursively create nested folders
            create_random_folder_structure(folder_path, max_depth, current_depth + 1, uid, gid)

            # Set ownership and permissions for the folder (rwxr-xr-x for folders)
            os.chmod(folder_path, 0o755)  # Fixed permissions for folders
            set_ownership(folder_path, uid, gid)  # Set ownership for folders to user
            print(f'Set ownership and permissions for directory: {folder_path}')

# Main function
if __name__ == "__main__":
    base_directory = "./random_structure"
    os.makedirs(base_directory, exist_ok=True)

    # Get ownership information of the base directory
    stat_info = os.stat(base_directory)
    uid = stat_info.st_uid  # User ID
    gid = stat_info.st_gid  # Group ID

    # Get root user and group IDs
    root_uid = 0  # Root user ID
    root_gid = 0  # Root group ID

    # Fetch username and group name for verification (optional)
    username = pwd.getpwuid(uid).pw_name
    groupname = grp.getgrgid(gid).gr_name
    print(f'Creating folder structure as user: {username}, group: {groupname}')

    # Create the random folder structure
    create_random_folder_structure(base_directory, uid=uid, gid=gid)

    # Randomly assign ownership to root after all files and directories are created
#    randomly_assign_root_ownership(base_directory, uid, gid, root_uid, root_gid)

    print(f'Random folder structure created at {os.path.abspath(base_directory)}')