#!/usr/bin/env python# -*- coding: utf-8 -*--# Copyright (c) 2022 Oracle and/or its affiliates.# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/importosimportsubprocessfromcontextlibimportcontextmanagerDEFAULT_KRB5_CONFIG_PATH="~/.bds_config/krb5.conf"KRB5_CONFIG="KRB5_CONFIG"
[docs]classKRB5KinitError(Exception):# pragma: no cover"""KRB5KinitError class when kinit -kt command failed to generate cached ticket with the keytab file and the krb5 config file."""pass
[docs]definit_ccache_with_keytab(principal:str,keytab_file:str)->None:"""Initialize credential cache using keytab file. Parameters ---------- principal: str The unique identity to which Kerberos can assign tickets. keytab_path: str Path to your keytab file. Returns ------- None Nothing. """cmd="kinit -kt %(keytab_file)s%(principal)s"args={}args["principal"]=principalargs["keytab_file"]=keytab_filekinit_proc=subprocess.Popen((cmd%args).split(),stderr=subprocess.PIPE)stdout_data,stderr_data=kinit_proc.communicate()ifkinit_proc.returncode>0:raiseKRB5KinitError(stderr_data)
[docs]@contextmanagerdefkrbcontext(principal:str,keytab_path:str,kerb5_path:str=DEFAULT_KRB5_CONFIG_PATH)->None:"""A context manager for Kerberos-related actions. It provides a Kerberos context that you can put code inside. It will initialize credential cache automatically with keytab if no cached ticket exists. Otherwise, does nothing. Parameters ---------- principal: str The unique identity to which Kerberos can assign tickets. keytab_path: str Path to your keytab file. kerb5_path: (str, optional). Path to your krb5 config file. Returns ------- None Nothing. Examples -------- >>> from ads.bds.auth import krbcontext >>> from pyhive import hive >>> with krbcontext(principal = "your_principal", keytab_path = "your_keytab_path"): >>> hive_cursor = hive.connect(host="your_hive_host", ... port="your_hive_port", ... auth='KERBEROS', ... kerberos_service_name="hive").cursor() """refresh_ticket(principal=principal,keytab_path=keytab_path,kerb5_path=kerb5_path)yield
[docs]defrefresh_ticket(principal:str,keytab_path:str,kerb5_path:str=DEFAULT_KRB5_CONFIG_PATH)->None:"""generate new cached ticket based on the principal and keytab file path. Parameters ---------- principal: str The unique identity to which Kerberos can assign tickets. keytab_path: str Path to your keytab file. kerb5_path: (str, optional). Path to your krb5 config file. Returns ------- None Nothing. Examples -------- >>> from ads.bds.auth import refresh_ticket >>> from pyhive import hive >>> refresh_ticket(principal = "your_principal", keytab_path = "your_keytab_path") >>> hive_cursor = hive.connect(host="your_hive_host", ... port="your_hive_port", ... auth='KERBEROS', ... kerberos_service_name="hive").cursor() """keytab_path=os.path.abspath(os.path.expanduser(keytab_path))os.environ[KRB5_CONFIG]=os.path.abspath(os.path.expanduser(kerb5_path))ifnotos.path.exists(os.environ[KRB5_CONFIG]):raiseFileNotFoundError(f"krb5 config file not found in {kerb5_path}.")ifnothas_kerberos_ticket():init_ccache_with_keytab(principal,keytab_path)