codekingpro/portable-devtools

1##########################################################################2#3# pgAdmin 4 - PostgreSQL Tools4#5# Copyright (C) 2013 - 2024, The pgAdmin Development Team6# This software is released under the PostgreSQL Licence7#8#########################################################################9 10"""This file contains functions fetching different utility paths."""11 12import os13 14from flask import current_app, url_for15from flask_security import current_user16from werkzeug.exceptions import InternalServerError17from pgadmin.utils.constants import MY_STORAGE18from pgadmin.model import User19 20PGADMIN_PATH = '~/.pgadmin/'21 22 23def preprocess_username(un):24    ret_un = un25    if len(ret_un) == 0 or ret_un[0].isdigit():26        ret_un = 'pga_user_' + un27 28    ret_un = ret_un.replace('@', '_') \29        .replace('/', 'slash') \30        .replace('\\', 'slash')31 32    return ret_un33 34 35def get_storage_directory(user=current_user, shared_storage=''):36    # Don't move this import statement to the top of the file,37    # it throws circular import error.38    import config39    if config.SERVER_MODE is not True:40        return None41 42    is_shared_storage = False43    if shared_storage != MY_STORAGE and shared_storage:44        is_shared_storage = True45        selected_dir = [sdir for sdir in config.SHARED_STORAGE if46                        sdir['name'] == shared_storage]47        storage_dir = None48        if len(selected_dir) > 0:49            the_dir = selected_dir[0]['path']50            storage_dir = the_dir51    else:52        storage_dir = getattr(53            config, 'STORAGE_DIR',54            os.path.join(55                os.path.realpath(56                    os.path.expanduser(PGADMIN_PATH)57                ), 'storage'58            )59        )60 61    if storage_dir is None:62        return None63 64    username = preprocess_username(user.username.split('@')[0])65 66    # Figure out the old-style storage directory name67    old_storage_dir = os.path.join(68        storage_dir.decode('utf-8') if hasattr(storage_dir, 'decode')69        else storage_dir,70        username71    )72 73    username = preprocess_username(user.username)74 75    if is_shared_storage:76        # Figure out the new style storage directory name77        storage_dir = os.path.join(78            storage_dir.decode('utf-8') if hasattr(storage_dir, 'decode')79            else storage_dir80        )81    else:82        # Figure out the new style storage directory name83        storage_dir = os.path.join(84            storage_dir.decode('utf-8') if hasattr(storage_dir, 'decode')85            else storage_dir,86            username87        )88 89    # Rename an old-style storage directory, if the new style doesn't exist90    if os.path.exists(old_storage_dir) and not os.path.exists(storage_dir):91        current_app.logger.warning(92            'Renaming storage directory %s to %s.',93            old_storage_dir, storage_dir94        )95        os.rename(old_storage_dir, storage_dir)96 97    if not os.path.exists(storage_dir):98        os.makedirs(storage_dir, int('700', 8))99 100    return storage_dir101 102 103def init_app():104    # Don't move this import statement to the top of the file,105    # it throws circular import error.106    import config107    if config.SERVER_MODE is not True:108        return None109 110    storage_dir = getattr(111        config, 'STORAGE_DIR',112        os.path.join(113            os.path.realpath(114                os.path.expanduser(PGADMIN_PATH)115            ), 'storage'116        )117    )118 119    if storage_dir and not os.path.isdir(storage_dir):120        if os.path.exists(storage_dir):121            raise InternalServerError(122                'The path specified for the storage directory is not a '123                'directory.'124            )125        os.makedirs(storage_dir, int('700', 8))126 127    if storage_dir and not os.access(storage_dir, os.W_OK | os.R_OK):128        raise InternalServerError(129            'The user does not have permission to read and write to the '130            'specified storage directory.'131        )132 133 134def get_cookie_path():135    cookie_root_path = '/'136    pgadmin_root_path = url_for('browser.index')137    if pgadmin_root_path != '/browser/':138        cookie_root_path = pgadmin_root_path.replace(139            '/browser/', ''140        )141    return cookie_root_path142 143 144def create_users_storage_directory():145    """146    This function is used to iterate through all the users and147    create users directory if not already created.148    """149    # Don't move this import statement to the top of the file,150    # it throws circular import error.151    import config152    if not config.SERVER_MODE:153        return None154 155    users = User.query.all()156 157    for usr in users:158        username = preprocess_username(usr.username)159 160        storage_dir = getattr(161            config, 'STORAGE_DIR',162            os.path.join(163                os.path.realpath(164                    os.path.expanduser(PGADMIN_PATH)165                ), 'storage'166            )167        )168 169        if storage_dir is None:170            return None171 172        storage_dir = os.path.join(173            storage_dir.decode('utf-8') if hasattr(storage_dir, 'decode')174            else storage_dir, username175        )176 177        if not os.path.exists(storage_dir):178            os.makedirs(storage_dir, int('700', 8))179