Team Ai
Datasetpublic

codekingpro/portable-devtools

sourceHugging Faceupdated 5mo agoView on Hugging Face
1likes14kdownloads
utils.py293 linesDownload Raw Back to servers
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"""Server helper utilities"""11from ipaddress import ip_address12from werkzeug.exceptions import InternalServerError13from flask import render_template14 15from pgadmin.utils.crypto import encrypt, decrypt16import config17from pgadmin.model import db, Server18 19 20def is_valid_ipaddress(address):21    try:22        return bool(ip_address(address))23    except ValueError:24        return False25 26 27def parse_priv_from_db(db_privileges):28    """29    Common utility function to parse privileges retrieved from database.30    """31    acl = {32        'grantor': db_privileges['grantor'],33        'grantee': db_privileges['grantee'],34        'privileges': []35    }36    if 'acltype' in db_privileges:37        acl['acltype'] = db_privileges['acltype']38 39    privileges = []40    for idx, priv in enumerate(db_privileges['privileges']):41        privileges.append({42            "privilege_type": priv,43            "privilege": True,44            "with_grant": db_privileges['grantable'][idx]45        })46 47    acl['privileges'] = privileges48 49    return acl50 51 52def _check_privilege_type(priv):53    if isinstance(priv['privileges'], dict) \54            and 'changed' in priv['privileges']:55        tmp = []56        for p in priv['privileges']['changed']:57            tmp_p = {'privilege_type': p['privilege_type'],58                     'privilege': False,59                     'with_grant': False}60 61            if 'with_grant' in p:62                tmp_p['privilege'] = True63                tmp_p['with_grant'] = p['with_grant']64 65            if 'privilege' in p:66                tmp_p['privilege'] = p['privilege']67 68            tmp.append(tmp_p)69 70        priv['privileges'] = tmp71 72 73def _parse_privileges(priv, db_privileges, allowed_acls, priv_with_grant,74                      priv_without_grant):75    _check_privilege_type(priv)76    for privilege in priv['privileges']:77 78        if privilege['privilege_type'] not in db_privileges:79            continue80 81        if privilege['privilege_type'] not in allowed_acls:82            continue83 84        if privilege['with_grant']:85            priv_with_grant.append(86                db_privileges[privilege['privilege_type']]87            )88        elif privilege['privilege']:89            priv_without_grant.append(90                db_privileges[privilege['privilege_type']]91            )92 93 94def parse_priv_to_db(str_privileges, allowed_acls=[]):95    """96    Common utility function to parse privileges before sending to database.97    """98    from pgadmin.utils.driver import get_driver99    from config import PG_DEFAULT_DRIVER100    driver = get_driver(PG_DEFAULT_DRIVER)101 102    db_privileges = {103        'c': 'CONNECT',104        'C': 'CREATE',105        'T': 'TEMPORARY',106        'a': 'INSERT',107        'r': 'SELECT',108        'w': 'UPDATE',109        'd': 'DELETE',110        'D': 'TRUNCATE',111        'x': 'REFERENCES',112        't': 'TRIGGER',113        'U': 'USAGE',114        'X': 'EXECUTE'115    }116 117    privileges = []118    allowed_acls_len = len(allowed_acls)119 120    for priv in str_privileges:121        priv_with_grant = []122        priv_without_grant = []123 124        _parse_privileges(priv, db_privileges, allowed_acls, priv_with_grant,125                          priv_without_grant)126 127        # If we have all acl then just return all128        if len(priv_with_grant) == allowed_acls_len > 1:129            priv_with_grant = ['ALL']130        if len(priv_without_grant) == allowed_acls_len > 1:131            priv_without_grant = ['ALL']132 133        grantee = driver.qtIdent(None, priv['grantee']) \134            if priv['grantee'] != 'PUBLIC' else 'PUBLIC'135 136        old_grantee = driver.qtIdent(None, priv['old_grantee']) \137            if 'old_grantee' in priv and priv['old_grantee'] != 'PUBLIC' \138            else grantee139 140        acltype = priv['acltype'] if 'acltype' in priv else 'defaultacls'141 142        grantor = driver.qtIdent(None, priv['grantor'])143 144        # Appending and returning all ACL145        privileges.append({146            'grantor': grantor,147            'grantee': grantee,148            'with_grant': priv_with_grant,149            'without_grant': priv_without_grant,150            'old_grantee': old_grantee,151            'acltype': acltype152        })153 154    return privileges155 156 157def tokenize_options(options_from_db, option_name, option_value):158    """159    This function will tokenize the string stored in database160    e.g. database store the value as below161    key1=value1, key2=value2, key3=value3, ....162    This function will extract key and value from above string163 164    Args:165        options_from_db: Options from database166        option_name: Option Name167        option_value: Option Value168 169    Returns:170        Tokenized options171    """172    options = []173    if options_from_db is not None:174        for fdw_option in options_from_db:175            k, v = fdw_option.split('=', 1)176            options.append({option_name: k, option_value: v})177    return options178 179 180def validate_options(options, option_name, option_value):181    """182    This function will filter validated options183    and sets flag to use in sql template if there are any184    valid options185 186    Args:187        options: List of options188        option_name: Option Name189        option_value: Option Value190 191    Returns:192        Flag, Filtered options193    """194    valid_options = []195    is_valid_options = False196 197    for option in options:198        # If option name is valid199        if option_name in option and \200            option[option_name] is not None and \201                option[option_name] != '' and \202                len(option[option_name].strip()) > 0:203            # If option value is valid204            if option_value in option and \205                option[option_value] is not None and \206                option[option_value] != '' and \207                    len(option[option_value].strip()) > 0:208                # Do nothing here209                pass210            else:211                # Set empty string if no value provided212                option[option_value] = ''213            valid_options.append(option)214 215    if len(valid_options) > 0:216        is_valid_options = True217 218    return is_valid_options, valid_options219 220 221def _password_check(server, manager, old_key, new_key):222    # Check if old password was stored in pgadmin4 sqlite database.223    # If yes then update that password.224    if server.password is not None:225        password = decrypt(server.password, old_key)226 227        if isinstance(password, bytes):228            password = password.decode()229 230        password = encrypt(password, new_key)231        setattr(server, 'password', password)232        manager.password = password233 234 235def reencrpyt_server_passwords(user_id, old_key, new_key):236    """237    This function will decrypt the saved passwords in SQLite with old key238    and then encrypt with new key239    """240    from pgadmin.utils.driver import get_driver241    driver = get_driver(config.PG_DEFAULT_DRIVER)242 243    for server in Server.query.filter_by(user_id=user_id).all():244        manager = driver.connection_manager(server.id)245 246        _password_check(server, manager, old_key, new_key)247 248        if server.tunnel_password is not None:249            tunnel_password = decrypt(server.tunnel_password, old_key)250            if isinstance(tunnel_password, bytes):251                tunnel_password = tunnel_password.decode()252 253            tunnel_password = encrypt(tunnel_password, new_key)254            setattr(server, 'tunnel_password', tunnel_password)255            manager.tunnel_password = tunnel_password256        elif manager.tunnel_password is not None:257            tunnel_password = decrypt(manager.tunnel_password, old_key)258 259            if isinstance(tunnel_password, bytes):260                tunnel_password = tunnel_password.decode()261 262            tunnel_password = encrypt(tunnel_password, new_key)263            manager.tunnel_password = tunnel_password264 265        db.session.commit()266        manager.update_session()267 268 269def remove_saved_passwords(user_id):270    """271    This function will remove all the saved passwords for the server272    """273 274    try:275        db.session.query(Server) \276            .filter(Server.user_id == user_id) \277            .update({Server.password: None, Server.tunnel_password: None})278        db.session.commit()279    except Exception:280        db.session.rollback()281        raise282 283 284def get_replication_type(conn, sversion):285    status, res = conn.execute_dict(render_template(286        "/servers/sql/#{0}#/replication_type.sql".format(sversion)287    ))288 289    if not status:290        raise InternalServerError(res)291 292    return res['rows'][0]['type']293 
codekingpro/portable-devtools · Team Ai