
    `                        S r SSKJr  SSKJr  SSKJr  SSKrSSKrSSKrSSKrSSK	r	SSK
r
SSKrSSKrSSKJr  SSKJr  SSKJr  SSKJr  SS	KJr  SS
KJr  SSKJr  SSKJr  SSKJr  SSKJr  SSKJr  SSKJr  SSKJ r   SSK!J"r#  SSK!J$r%  SSK&J'r'  SSK(r(\RR                  RU                  \RR                  RW                  \,5      S5      r-S r.  S6S jr/S r0 " S S\15      r2S r3S r4S7S jr5S r6  S8S jr7 " S  S!\15      r8     S9S" jr9S# r:S$ r;S% r<S& r=S' r>S( r?S) r@S* rAS:S+ jrBS, rCS- rDS. rES/ rF " S0 S1\15      rGS:S2 jrHS3 rIS4 rJS5 rKg);z.Common utilities for the gcloud dataproc tool.    )absolute_import)division)unicode_literalsN)encoding)
exceptions)
list_pager)storage_helpers)arg_parsers)util)log)
properties)requests)console_attr)
console_io)progress_tracker)creds)store)retryschemasc                 t    [         R                  " S[        R                  " U 5      -   5        U R                  $ )zReturns a printable representation of a failed Google API's status.proto.

Args:
  error: the failed Status to print.

Returns:
  A ready-to-print string representation of the error.
zError:
)r   debugr   MessageToJsonmessageerrors    +lib/googlecloudsdk/api_lib/dataproc/util.pyFormatRpcErrorr   3   s*     ))J//667	    c                 R   [         R                  " USS9   [        R                  " 5       nU[        R                  " 5       U-
  :  a=   U " U5        [        R                  " U5        U[        R                  " 5       U-
  :  a  M=  SSS5        [        R                  " SR                  U5      5      e! [        R                   a     SSS5        g[        R
                   a2  n[        R                  " SX5        [        U5      (       a  e  SnANSnAff = f! , (       d  f       N= f)z1Poll Dataproc resource until it no longer exists.TautotickNz!Get request for [{0}] failed:
{1}z"Deleting resource [{0}] timed out.)r   ProgressTrackertimeapitools_exceptionsHttpNotFoundError	HttpErrorr   r   IsClientHttpExceptionsleepr   OperationTimeoutErrorformat)request_methodresource_refr   	timeout_spoll_period_s
start_timer   s          r   WaitForResourceDeletionr0   @   s     ''$?J
tyy{Z/
0
|$ jj tyy{Z/
0 @ 	((*11,?	A A !22  @? !** 		6L !''
 (	 @?sA   2D	B22D2DDD#(DDDD
D&c                  @    [         R                  " 5       R                  $ N)uuiduuid4hex r   r   GetUniqueIdr7   Y   s    			r   c                       \ rS rSrSrS rSrg)Bunch]   zwClass that converts a dictionary to javascript like object.

For example:
    Bunch({'a': {'b': {'c': 0}}}).a.b.c == 0
c                     [         R                  " U5       H3  u  p#[        U[        5      (       a  [	        U5      nX0R
                  U'   M5     g r2   )six	iteritems
isinstancedictr9   __dict__)self
dictionarykeyvalues       r   __init__Bunch.__init__d   s9    mmJ/
	E4	 	 e mmC 0r   r6   N)__name__
__module____qualname____firstlineno____doc__rE   __static_attributes__r6   r   r   r9   r9   ]   s    !r   r9   c                 H    U R                  SSSS9  U R                  SSSS9  g )Nz--jarmain_jarz3The HCFS URI of jar file containing the driver jar.)desthelpz--class
main_classzuThe class containing the main method of the driver. Must be in a provided jar or jar that is already on the classpath)add_argument)parsers    r   AddJvmDriverFlagsrT   k   sA    @  B 	D  Fr   c                 L    U R                   S:  =(       a    U R                   S:  $ )z>Returns true if the http exception given is an HTTP 4xx error.i  i  )status_code)http_exceptions    r   r'   r'   w   s#    		#	#s	*	O~/I/IC/OOr   c                   ^^ U R                   R                  UR                  S9n[        R                  R                  SR                  UR                  5      5        [        R                  " 5       nSm[        R                  " SS9nU(       a  SOSmUU4S jn[        R                  " USS	9   U[        R                  " 5       U-
  :  a   U R                  R                  R                  U5      n[        UR                   U R                   R"                  5      n	U" U	R$                  5        ['        U	R$                  5      mUR(                  (       a  O5 [        R0                  " U5        U[        R                  " 5       U-
  :  a  M  S
S
S
5        [        UR                   U R                   R"                  5      n	U" U	R$                  5        UR(                  (       d/  [2        R4                  " SR                  UR                  5      5      eUR6                  (       aC  [2        R8                  " SR                  UR                  [;        UR6                  5      5      5      e[        R<                  " SUR                  [        R                  " 5       U-
  5        U$ ! [*        R,                   a  n
[/        U
5      (       a  e  S
n
A
GNqS
n
A
ff = f! , (       d  f       GNL= f)a1  Poll dataproc Operation until its status is done or timeout reached.

Args:
  dataproc: wrapper for Dataproc messages, resources, and client
  operation: Operation, message of the operation to be polled.
  message: str, message to display to user while polling.
  timeout_s: number, seconds to poll with retries before timing out.
  poll_period_s: number, delay in seconds between requests.

Returns:
  Operation: the return value of the last successful operations.get
  request.

Raises:
  OperationError: if the operation times out or finishes with an error.
nameWaiting on operation [{0}].r   Tr   
 c                    > U TS  nU(       a?  [         R                  R                  T5        U H  n[         R                  " U5        M     g g r2   )r   errwritewarning)warningsnew_warningsra   tracker_separatorwarnings_so_fars      r   _LogWarnings&WaitForOperation.<locals>._LogWarnings   s?    O,-L	ggmm%&!'G " r   r    NOperation [{0}] timed out.Operation [{0}] failed: {1}.*Operation [%s] finished after %.3f seconds)messages+DataprocProjectsRegionsOperationsGetRequestrZ   r   statusPrintr*   r#   r   IsInteractiver   r"   clientprojects_regions_operationsGetParseOperationJsonMetadatametadataClusterOperationMetadatarb   lendoner$   r&   r'   r(   r   r)   r   OperationErrorr   info)dataproc	operationr   r-   r.   requestr/   is_ttyrf   rt   rW   rd   re   s              @@r   WaitForOperationr~   }   s0   " II>> J '**077	GHyy{*/##$/&$d" ''$?
tyy{Z/
0OO??CCGL	- 1 1 J JLX&&'h//0>>
  jj tyy{Z/
0 @ ((++DDF(x  !	

*
*+G+N+N,  

#
#$B$I$Iy7%9 : : ((7IIK*$'	# !**  00
 1 @?s7   0KBJ4KK1KKKK
Kc                 	   SnXA;  d  U R                   X   :w  a  U R                  bD  [        R                  R	                  SR                  U R                  U R                   5      5        O8[        R                  R	                  SR                  U R                   5      5        U R                   X'   U R                  US   :w  GaT  [        U R                  S5      (       aK  U R                  R                  b4  [        R                  R	                  U R                  R                  5        O[        U R                  S5      (       aP  U R                  R                  b9  [        R                  R	                  SR                  U R                  5      5        Ot[        U R                  S	5      (       aY  U R                  R                  bB  [        R                  R	                  S
R                  U R                  R                  5      5        U R                  US'   [        U R                  S5      (       GaU  U R                  R                   GH:  nUR                  (       d  M  UR                  U;  d  XR                     UR                   :w  a[  [        R                  R	                  SR                  UR                  UR                   5      5        UR                   XR                  '   UR                  (       d  M  UR                  U;  d  X5R                     UR                  :w  d  M  [        R                  R	                  SR                  UR                  UR                  5      5        UR                  X5R                  '   GM=     U R                  US   :w  GaU  [        U R                  S5      (       aK  U R                  R                  b4  [        R                  R	                  U R                  R                  5        O[        U R                  S5      (       aP  U R                  R                  b9  [        R                  R	                  SR                  U R                  5      5        Ot[        U R                  S	5      (       aY  U R                  R                  bB  [        R                  R	                  SR                  U R                  R                  5      5        U R                  US'   gg)aD  Print workflow and job status for the running workflow template.

This method will detect any changes of state in the latest metadata and print
all the new states in a workflow template.

For example:
  Workflow template template-name RUNNING
  Creating cluster: Operation ID create-id.
  Job ID job-id-1 RUNNING
  Job ID job-id-1 COMPLETED
  Deleting cluster: Operation ID delete-id.
  Workflow template template-name DONE

Args:
  metadata: Dataproc WorkflowMetadata message object, contains the latest
    states of a workflow template.
  status: Dictionary, stores all jobs' status in the current workflow
    template, as well as the status of the overarching workflow.
  operations: Dictionary, stores cluster operation status for the workflow
    template.
  errors: Dictionary, stores errors from the current workflow template.
wtNzWorkflowTemplate [{0}] {1}zWorkflowTemplate {0}createClusterr   rw   zCreated cluster: {0}.operationIdz%Creating cluster: Operation ID [{0}].nodeszJob ID {0} {1}zJob ID {0} error: {1}deleteClusterzDeleted cluster: {0}.z%Deleting cluster: Operation ID [{0}].)statetemplater   rm   rn   r*   r   hasattrr   rw   clusterNamer   graphr   jobIdr   )rt   rm   
operationserrorstemplate_keynodes         r   PrintWorkflowMetadatar      s   0 ,8>>V5I#I$	jj3::


X^^- . 
jj-44X^^DE#>>Fz/::x%% $2288D	jjx--334	''
 
%3388D	jj.55h6J6JKL	
 
#11==I	jj>EE

 
 
,
,. /"*"8"8JX^^W%%$$ZZ	6	!VJJ%74::%E

)00TZZHI!ZZzz	61

+tzz9

077

DJJOP!ZZzz % z/::x%% $2288D	jjx--334	''
 
%3388D	jj.55h6J6JKL	
 
#11==I	jj>EE

 
 
,
,. /"*"8"8J ;r   c                    U R                   R                  UR                  S9n[        R                  R                  SR                  UR                  5      5        [        R                  " 5       nSSS.n0 n0 nUb  U[        R                  " 5       U-
  :  a   U R                  R                  R                  U5      n[        UR                  U R                   R                  5      n	[        XXh5        UR                  (       a  O: [        R&                  " U5        Uc  M  U[        R                  " 5       U-
  :  a  M  [        UR                  U R                   R                  5      n	UR                  (       d/  [(        R*                  " SR                  UR                  5      5      eUR,                  (       aC  [(        R.                  " SR                  UR                  [1        UR,                  5      5      5      eS Hb  nX;   d  M
  Xk   c  M  Xk   R,                  (       d  M&  [(        R.                  " SR                  Xk   R2                  Xk   R,                  5      5      e   [        R4                  " SUR                  [        R                  " 5       U-
  5        U$ ! [         R"                   a  n
[%        U
5      (       a  e  Sn
A
GNSn
A
ff = f)a  Poll dataproc Operation until its status is done or timeout reached.

Args:
  dataproc: wrapper for Dataproc messages, resources, and client
  operation: Operation, message of the operation to be polled.
  timeout_s: number, seconds to poll with retries before timing out.
  poll_period_s: number, delay in seconds between requests.

Returns:
  Operation: the return value of the last successful operations.get
  request.

Raises:
  OperationError: if the operation times out or finishes with an error.
rY   r[   N)r   r   rh   ri   rj   )rk   rl   rZ   r   rm   rn   r*   r#   rp   rq   rr   rs   rt   WorkflowMetadatar   rw   r$   r&   r'   r(   r   r)   r   rx   r   r   ry   )rz   r{   r-   r.   r|   r/   r   rm   r   rt   rW   ops               r    WaitForWorkflowTemplateOperationr     sE   & II>> J '**077	GHyy{*!%=*&& 	Y$))+
*BC//==AA'Ji+I,>,>,4,=,=,N,NPh HjA	 
 	JJ} 	Y$))+
*BC (	(:(:(0(9(9(J(JL( 


*
*+G+N+N,  

#
#$B$I$Iy7%9 : :.b	JN6:>;O;O;O%%&D&K&K
.
$
$jn&:&:'< = = /
 ((7IIK*$'	+ (( 	~	.	. 
/s   A,J! !K5KKc                   $    \ rS rSrSrS rS rSrg)NoOpProgressDisplayiB  z8For use in place of a ProgressTracker in a 'with' block.c                     g r2   r6   rA   s    r   	__enter__NoOpProgressDisplay.__enter__E      r   c                     g r2   r6   )rA   unused_argss     r   __exit__NoOpProgressDisplay.__exit__H  r   r   r6   N)rG   rH   rI   rJ   rK   r   r   rL   r6   r   r   r   r   B  s    @		r   r   c
           	      B	  ^ U R                   R                  UR                  UR                  UR                  S9n
SmSnSnSnSnU4S jnS nU(       a*  [
        R                  R                  S5        [        5       nO[        R                  " USS	9n[        R                  " 5       =nnU   U	(       a
  U	UU-
  :  Ga  U" 5         T=(       a    TR                  (       + nU(       d.  UR                  R                  U R                  ;   a
  SnUU-
  S
-   n	U(       a  U(       a  U(       a  GO2U(       + =(       a    UX-   :  nU=(       a    T(       + nU(       + =(       a    UnU(       d  U(       d  U(       a  Un U R                  R                   R#                  U
5      nU(       ax  UR2                  (       ag  UR2                  U:w  aW  U(       a$  U" 5         [
        R(                  " S5        U" 5         UR2                  n[4        R6                  " UR2                  5      m[        R8                  " U5        [        R                  " 5       nU	(       d  GM  U	UU-
  :  a  GM  SSS5        UR                  R                  nUU R                  ;   Gaf  U(       aE  T(       d  [
        R(                  " S5        O'TR                  (       a  [
        R(                  " S5        UUL a  U$ U(       a  UUL a  UR                  R:                  (       aD  [<        R>                  " SR+                  UR                  UR                  R:                  5      5      e[<        R>                  " SR+                  UR                  5      5      eUR                  R:                  (       a-  [
        R@                  " SUR                  R:                  -   5        [<        R>                  " SR+                  UR                  UU5      5      e[<        RB                  " SR+                  UR                  U5      5      e! [$        R&                   aU  n[
        R(                  " SR+                  [,        R.                  " U5      5      5        [1        U5      (       a  e  SnAGNSnAff = f! , (       d  f       GN8= f)a  Poll dataproc Job until its status is terminal or timeout reached.

Args:
  dataproc: wrapper for dataproc resources, client and messages
  job: The job to wait to finish.
  job_ref: Parsed dataproc.projects.regions.jobs resource containing a
    projectId, region, and jobId.
  message: str, message to display to user while polling.
  goal_state: JobStatus.StateValueValuesEnum, the state to define success
  error_state: JobStatus.StateValueValuesEnum, the state to define failure
  stream_driver_log: bool, Whether to show the Job's driver's output.
  log_poll_period_s: number, delay in seconds between checking on the log.
  dataproc_poll_period_s: number, delay in seconds between requests to the
    Dataproc API.
  timeout_s: number, time out for job completion. None means no timeout.

Returns:
  Job: the return value of the last successful jobs.get request.

Raises:
  JobError: if the job finishes with an error.
)	projectIdregionr   Nr   Fc                  x   > T (       a2  T R                   (       a   T R                  [        R                  5        g g g r2   )openReadIntoWritabler   r_   )driver_log_streams   r   ReadDriverLogIfPresent5WaitForJobTermination.<locals>.ReadDriverLogIfPresentt  s)    .33((1 4r   c                      [         R                  " 5       n [        R                  R	                  SU R                  5       S   -  5        g )N=r   )r   GetConsoleAttrr   r_   rn   GetTermSize)attrs    r   PrintEqualsLine.WaitForJobTermination.<locals>.PrintEqualsLiney  s3    &&(DGGMM#((*1--.r   zWaiting for job output...Tr    
   zGetJob failed:
{}z3Job attempt failed. Streaming new attempt's output.zExpected job output not found.z4Job terminated, but output did not finish streaming.z Job [{0}] failed with error:
{1}zJob [{0}] failed.z	Details:
z6Job [{0}] entered state [{1}] while waiting for [{2}].z)Job [{0}] timed out while in state [{1}].)"rk   %DataprocProjectsRegionsJobsGetRequestr   r   r   r   rm   rn   r   r   r"   r#   r   r   terminal_job_statesrp   projects_regions_jobsrr   r$   r&   ra   r*   r<   	text_typer'   driverOutputResourceUrir	   StorageObjectSeriesStreamr(   detailsr   JobErrorry   JobTimeoutError)rz   jobjob_refr   
goal_stateerror_statestream_driver_loglog_poll_period_sdataproc_poll_period_sr-   r|   last_job_poll_timejob_completewait_displaydriver_output_urir   r   r/   nowlog_stream_closedregular_job_pollexpecting_output_streamexpecting_job_doner   r   r   s                            @r   WaitForJobTerminationr   L  s   @ CC!!'.. D O',,2
/ JJ01&(L#33GdKLYY[ *s9j(89+J4E4J4J0J
**

h::
:*$r)		04E 
 A'@@  !2 K:K6K++A0A	48J 	5599'B# #"="=''+<<KKMN!99
-GG))+

jj"#IIKcS i9j(89 X **

% h***45!!JK
ju+			!!"E"L"LMM3::--#/ 0 	0 3 : :7== IJJ
zz	hh|cjj0001


@GGMM5*	./ / 	""188N	P PM #,, 	
++*11#--2FG
H"5)) *	7 |s?   -CR9%P#B1R	R#R7A
RRRR
Rc                  f    [         R                  R                  R                  R	                  5       $ r2   )r   VALUESrz   r   	GetOrFailr6   r   r   ResolveRegionr     s#    				#	#	*	*	4	4	66r   c                  f    [         R                  R                  R                  R	                  5       $ r2   )r   r   rz   locationr   r6   r   r   ResolveLocationr     s#    				#	#	,	,	6	6	88r   c                     UR                   R                  U [        [        R                  R
                  R                  R                  S.SS9nU$ )Nr   r   z"dataproc.projects.regions.clustersparams
collection	resourcesParser   r   r   coreprojectr   )rZ   rz   refs      r   ParseClusterr     sO      
!!((--55?? 6 	! 	7# 
*r   c                     UR                   R                  U [        [        R                  R
                  R                  R                  S.SS9nU$ )Nr   zdataproc.projects.regions.jobsr   r   )job_idrz   r   s      r   ParseJobr     sO      !!((--55?? 2 	! 	3# 
*r   c                 t    U (       d  U" 5       $ [         R                  " U[         R                  " U 5      5      $ )z2Returns an Operation message for a metadata value.)r   JsonToMessager   )metadata_valuemetadata_types     r   rs   rs     s2    	?			 ( 6 6~ F
H Hr   c                     U R                   R                  S [        [        R                  R
                  R                  R                  S.SS9nU$ )N)regionIdr   zdataproc.projects.regionsr   r   rz   r   s     r   ParseRegionr      sO      
#!((--55?? - 	! 	.# 
*r   c                     U R                   R                  S [        [        R                  R
                  R                  R                  S.SS9nU$ N)locationsId
projectsIdzdataproc.projects.locationsr   r   r   s     r   ParseProjectsLocationsr     sO      
&"))..66@@ / 	! 	0# 
*r   c                     U R                   R                  S [        5       [        R                  R
                  R                  R                  S.SS9nU$ r   )r   r   r   r   r   r   r   r   r   s     r    ParseProjectsLocationsForSessionr     sR      
(*"))..66@@ / 	! 	0# 
*r   c                    [         R                  " U=(       d    SSS9n[        R                  " U R                  R
                  US9nXl        SUl        UR                  b  UR                  R                  bJ  [        [        R                  " SSS9" UR                  R                  5      5      S	-   UR                  l	        UR                  R                  R                  b^  [        [        R                  " S
SS9" UR                  R                  R                  5      5      S	-   UR                  R                  l        U$ )a{  Returns autoscaling policy read from YAML file.

Args:
  dataproc: wrapper for dataproc resources, client and messages.
  policy_id: The autoscaling policy id (last piece of the resource name).
  policy_file_name: if set, location of the YAML file to read from. Otherwise,
    reads from stdin.

Raises:
  argparse.ArgumentError if duration formats are invalid or out of bounds.
-Fbinarymessage_typestreamN2m1d)lower_boundupper_bounds0s)r   ReadFromFileOrStdinexport_utilImportrk   AutoscalingPolicyidrZ   basicAlgorithmcooldownPeriodstrr
   Duration
yarnConfiggracefulDecommissionTimeout)rz   	policy_idpolicy_file_namedatapolicys        r   ReadAutoscalingPolicyr  %  s    
	'	'(8(?C	N$$$66tE&
 ) &+ &++7-0


4T
B##224.57:.;f* ''CCOEH


4T
B  ++GGIFJLOFPf&&B 
-r   c                 :   SR                  UR                  S5      SS 5      nU R                  R                  UUS9nU R                  R
                  R                  U5      n[        R                  R                  SR                  UR                  5      5        U$ )zReturns the server-resolved policy after creating the given policy.

Args:
  dataproc: wrapper for dataproc resources, client and messages.
  name: The autoscaling policy resource name.
  policy: The AutoscalingPolicy message to create.
/r      )parentautoscalingPolicyCreated [{0}].)joinsplitrk   7DataprocProjectsRegionsAutoscalingPoliciesCreateRequestrp   $projects_regions_autoscalingPoliciesCreater   rm   rn   r*   r   )rz   rZ   r
  r  r|   s        r   CreateAutoscalingPolicyr  K  s     88DJJsOAa()& MM  N " 
 ????FFwO&**#**69956	-r   c                     Xl         U R                  R                  R                  U5      n[        R
                  R                  SR                  UR                  5      5        U$ )zReturns the server-resolved policy after updating the given policy.

Args:
  dataproc: wrapper for dataproc resources, client and messages.
  name: The autoscaling policy resource name.
  policy: The AutoscalingPolicy message to create.
Updated [{0}].)	rZ   rp   r  Updater   rm   rn   r*   r   )rz   rZ   r
  s      r   UpdateAutoscalingPolicyr  b  sO     + OO88??G 	**#**69956	-r   c                    SSSU US.n[         R                  R                  R                  R	                  5       nSU S3n[         R                  R
                  R                  R                  5       (       a  SU S3nSS0n[        R                  " 5       R                  XEUS	9nUR                  S
:w  a  [        S5      e[        R                  " UR                  5      nUR!                  SS5      $ )zDownscope the given credentials to the given access boundary.

Args:
  token: The credentials to downscope.
  access_boundary_json: The JSON-formatted access boundary.

Returns:
  A downscopded credential with the given access-boundary.
z/urn:ietf:params:oauth:grant-type:token-exchangez-urn:ietf:params:oauth:token-type:access_token)
grant_typerequested_token_typesubject_token_typesubject_tokenoptionszhttps://sts.z	/v1/tokenzhttps://sts.mtls.zContent-Typez!application/x-www-form-urlencoded)headersr	     zError downscoping credentialsaccess_tokenN)r   r   r   universe_domainrr   context_awareuse_client_certificateGetBoolr   
GetSessionpostrV   
ValueErrorjsonloadscontentget)tokenaccess_boundary_jsonpayloadr%  cab_token_urlr"  downscope_response	cab_tokens           r   _DownscopeCredentialsr6  t  s     FMK%' %%**::>>@/  1;-$$;;CCEE''8	BM@A'**,117 2 4##s*
4
55jj+334)	~t	,,r   c                    [         R                  " SSSS9n[         R                  " U5        [        R                  " U5      (       a  UR
                  nOUR                  nU(       d  [        R                  " S5      e[        X 5      $ )a,  Get an access token for the user's current credentials.

Args:
  access_boundary_json: JSON string holding the definition of the access
    boundary to apply to the credentials.

Raises:
  PersonalAuthError: If no access token could be fetched for the user.

Returns:
  An access token for the user.
NT)allow_account_impersonationuse_google_authz?No access token could be obtained from the current credentials.)
c_storeLoadRefreshc_credsIsOauth2ClientCredentialsr$  r0  r   PersonalAuthErrorr6  )r1  credr0  s      r   GetCredentialsrA    sq     

d
D$	//$&&t,,EJJE	

&
&IK K	u	;;r   c                   l    \ rS rSrSrS rSS jrS rS r\	R                  " SS	9S
 5       rS rS rSrg)PersonalAuthUtilsi  z2Util functions for enabling personal auth session.c                     g r2   r6   r   s    r   rE   PersonalAuthUtils.__init__  r   r   Nc                    U/nUR                  U5        Sn [        [        SS5      (       ay  [        R                  " UU[        R                  [        R                  SS9nUR
                  R                  S5      R                  5       nUR                  5         UR                  $ [        R                  " U[        R                  [        R                  [        R                  S9nUR                  US9u  pU$ ! [         au  n
U(       a'  [        R                  " SS	R                  U5      U5        [         R"                  " S
S	R                  U5      -   S-   [$        R&                  " U
5      -   5      eSn
A
ff = f)aE  Run the specified command, capturing and returning output as appropriate.

Args:
  openssl_executable: The path to the openssl executable.
  args: The arguments to the openssl command to run.
  stdin: The input to the command.

Returns:
  The output of the command.

Raises:
  PersonalAuthError: If the call to openssl fails
NrunF)inputstdoutstderrcheckutf-8)stdinrI  rJ  )rH  z3OpenSSL command "%s" failed with error message "%s" z"Failure running openssl command: "z": )extendgetattr
subprocessrG  PIPErJ  decodestripcheck_returncoderI  Popencommunicate	Exceptionr   r   r  r   r?  r<   r   )rA   openssl_executableargsrM  commandrJ  procprI  _exs              r   _RunOpensslCommand$PersonalAuthUtils._RunOpensslCommand  s?    ""GNN4F!<	UD	)	)~~???? ##G,224 	{{//????	$
 MMM.	 <			G((7#V	-(()M),'):*;=B*C),r):*; < <	<s    BC; 'AC; ;
E:A0E55E:c                    U R                  USSSU/US9R                  S5      n UR                  5       R                  S5      S   n[	        U5      S:w  a  [        S	5      e[        US
5        US-   R                  S5      $ ! [         a2  n[        R                  " S[        R                  " U5      -   5      eSnAff = f)zCompute HMAC tag using OpenSSL.dgstz-sha256z-hmacrM  rL  rN     @   z1HMAC output is expected to be 64 characters long.   z'Failure due to invalid openssl output: Nr\   )r`  rS  rT  r  rv   r+  intrX  r   r?  r<   r   encode)rA   rC   r	  rY  
cmd_outputstripped_outputr_  s          r   _ComputeHmacPersonalAuthUtils._ComputeHmac  s    ((VY= ) F7O I"((*005a8o	_		#LMM	/2 d"**733  I((
3cmmB6G
GI IIs   AB 
B>-B99B>c                 $   [        U5      S:w  a  [        S5      e[        R                  " U5      R	                  S5      nU R                  USU5      n[        U5      nUR                  U5        UR                  S5        U R                  XFU5      $ )a]  Derives HMAC-based Key Derivation Function (HKDF) key through expansion on the initial pseudorandom key.

Args:
  prk: a pseudorandom key.
  info: optional context and application specific information (can be
    empty).
  openssl_executable: The path to the openssl executable.

Returns:
  Output keying material, expected to be of 256-bit length.
    zCThe given initial pseudorandom key is expected to be 32 bytes long.rL  r      )rv   r+  base64	b16encoderS  rl  	bytearrayrO  )rA   prkry   rY  
base16_prkt1t2datas          r   _DeriveHkdfKey PersonalAuthUtils._DeriveHkdfKey  s     3x2~
OQ Q!!#&--g6J			:s,>	?Br]F
MM$
MM'Z1CDDr   re  )max_retrialsc           
         [         R                  " US-   R                  S5      5      R                  5       n[        R
                  " [        R                  " S5      5      nUR                  S5      n[        R                  " S5      nU R                  USR                  S5      U5      n[        R
                  " U R                  USR                  S5      U5      5      R                  S5      n	[        R                  " 5        n
U
R                  UR                  S5      5        U
R                  S5        U R                  USS	S
SSU
R                  /[        R                   " U5      S9nSSS5        [#        W5      S:w  a  [%        S5      e[        R                   " U5      R                  S5      n[        R                  " 5        nUR                  U5        UR                  S5        SSSSUSSR'                  UR                  5      /nU R                  X>UR                  S5      S9nSSS5        [#        U5      S:w  a  [%        S5      e[        R                   " W5      R                  S5      n[)        U5      nUR+                  U5        U R-                  U	UU5      R                  S5      SS nSR'                  UUUUU5      $ ! , (       d  f       GNd= f! , (       d  f       N= f)zEncode token using OpenSSL.

Args:
  public_key: The public key for the session/cluster.
  secret: Token to be encrypted.
  openssl_executable: The path to the openssl executable.

Returns:
  Encrypted token.
r\   rL  rg  ro  encryption_keyauth_keyr   rsautlz-oaepz-encryptz-pubinz-inkeyrd  Ni   z3The encrypted key is expected to be 512 bytes long.encz-aes-256-ctrz-saltz-ivz-passzfile:{}z{}:{}:{}:{}:{})hashlibsha256ri  	hexdigestrq  rr  osurandomrS  rx  tempfileNamedTemporaryFiler`   seekr`  rZ   	b64encoderv   r+  r*   rs  rO  rl  )rA   
public_keysecretrY  key_hashiv_bytesinitialization_vectorinitial_keyr|  r}  kfencrypted_keyencoded_keypfencrypt_argsencrypted_tokenencoded_token
hmac_inputhmac_tags                      r   _EncodeTokenUsingOpenssl*PersonalAuthUtils._EncodeTokenUsingOpenssl  s    ~~zD088ABLLNH

2/H$OOG4**R.K(()9)@)@)I);=N K):):7)C.	0117  
	$	$	&"hhz  )*ggaj--
Wj(Hbgg
F  - . /m 
' =S LMM""=188AK		$	$	&"hh~ggaj
%1F


277
#l //
&--2H 0 Jo 
' =S LMM$$_5<<WEM8$Jo&  :!355;VG_%&r+H ""8]K#8(D D; 
'	& 
'	&s   A'KA$K
K
K(c                    U(       a  U R                  XU5      $  SSKnSSKJn  UR                  5         Sn[        R                  " U5      S   S   S   S   n[        R                  " US	-   R                  S
5      5      R                  5       nUR                  U5      n	UR                  U	5      n
U
R                  UR                   5      nUR#                  UR                  S
5      U5      n[$        R&                  " U5      R)                  S
5      nSR+                  X5      $ ! [         a    [        R
                  " S5      ef = f)zEncrypt secret with resource public key.

Args:
  public_key: The public key for the session/cluster.
  secret: Token to be encrypted.
  openssl_executable: The path to the openssl executable.

Returns:
  Encrypted token.
r   NhybridaC  Cannot load the Tink cryptography library. Either the library is not installed, or site packages are not enabled for the Google Cloud SDK. Please consult Cloud Dataproc Personal Auth documentation on adding Tink to Google Cloud SDK for further instructions.
https://cloud.google.com/dataproc/docs/concepts/iam/personal-authr   rC   keyDatarD   r\   rL  z{}:{})r  tinkr  ImportErrorr   r?  registerr,  r-  r  r  ri  r  JsonKeysetReaderread_no_secret_keyset_handle	primitiveHybridEncryptencryptrq  r  rS  r*   )rA   r  r  rY  r  r  contextpublic_key_valuer  readerkh_pub	encrypter
ciphertextr  s                 r   EncryptWithPublicKey&PersonalAuthUtils.EncryptWithPublicKeyI  s9    **:+=? ?O OOG zz*-e4Q7	B7K~~	D	 ((133<9;  "":.F..v6F   !5!56I""6==#97CJ$$Z077@M>>(223  O((NO OOs   
D# #!Ec                 :     SSK nSSK Jn  g! [         a     gf = f)z1Check if Tink cryptography library can be loaded.r   Nr  TF)r  r  r  )rA   r  r  s      r   IsTinkLibraryInstalled(PersonalAuthUtils.IsTinkLibraryInstalledw  s&    	   s   
 
r6   r2   )rG   rH   rI   rJ   rK   rE   r`  rl  rx  r   RetryOnExceptionr  r  r  rL   r6   r   r   rC  rC    sH    :	2<h4 E: 	q)3D *3Dj,3\r   rC  c                     [         R                  " U=(       d    SSS9n[        R                  " U R                  R
                  US9nU$ )a1  Returns session template read from YAML file.

Args:
  dataproc: Wrapper for dataproc resources, client and messages.
  template_file_name: If set, location of the YAML file to read from.
    Otherwise, reads from stdin.

Raises:
  argparse.ArgumentError if duration formats are invalid or out of bounds.
r   Fr   r   )r   r   r   r   rk   SessionTemplate)rz   template_file_namer	  r   s       r   ReadSessionTemplater    sE     
	'	'(:(Ac%	P$$$44TC( 
/r   c                 F   SR                  UR                  S5      SS 5      nXl        U R                  R	                  UUS9nU R
                  R                  R                  U5      n[        R                  R                  SR                  UR                  5      5        U$ )zReturns the server-resolved template after creating the given template.

Args:
  dataproc: Wrapper for dataproc resources, client and messages.
  name: The session template resource name.
  template: The SessionTemplate message to create.
r  r   r  )r  sessionTemplater  )r  r  rZ   rk   6DataprocProjectsLocationsSessionTemplatesCreateRequestrp   #projects_locations_sessionTemplatesr  r   rm   rn   r*   )rz   rZ   r   r  r|   s        r   CreateSessionTemplater    s     88DJJsOAa()&- NN" O $ 
 __@@GGP(**#**8==9:	/r   c                     Xl         U R                  R                  R                  U5      n[        R
                  R                  SR                  UR                   5      5        U$ )zReturns the server-resolved template after updating the given template.

Args:
  dataproc: Wrapper for dataproc resources, client and messages.
  name: The session template resource name.
  template: The SessionTemplate message to create.
r  )rZ   rp   r  Patchr   rm   rn   r*   )rz   rZ   r   s      r   UpdateSessionTemplater    sJ     -__@@FFxP(**#**8==9:	/r   c              /      ^#    [        5       mU4S jn[        R                  " USU0UD6nU H  nUv   M	     T(       a0  [        R                  " U SR                  [        T5      5      5        gg7f)z7Yields from paged List calls handling unreachable list.c                 P   > TR                  U R                  5        [        X5      $ r2   )updateunreachablerP  )r   r   r  s     r   _GetFieldFn5YieldFromListWithUnreachableList.<locals>._GetFieldFn  s"    w**+7!!r   get_field_funcz, N)setr   YieldFromListr   ra   r  sorted)unreachable_warning_msgrZ  kwargsr  resultitemr  s         @r    YieldFromListWithUnreachableListr    se     +" ##P;PP&d
J KK		&%& s   A.A1)<      )r  )Nr  )NFre  r   Nr2   )LrK   
__future__r   r   r   rq  r  r,  r  rQ  r  r#   r3   apitools.base.pyr   r   r$   r   googlecloudsdk.api_lib.dataprocr	   googlecloudsdk.callioper
   !googlecloudsdk.command_lib.exportr   r   googlecloudsdk.corer   r   r   googlecloudsdk.core.consoler   r   r   googlecloudsdk.core.credentialsr   r=  r   r:  googlecloudsdk.core.utilr   r<   pathr  dirname__file__
SCHEMA_DIRr   r0   r7   objectr9   rT   r'   r~   r   r   r   r   r   r   r   r   rs   r   r   r   r  r  r  r6  rA  rC  r  r  r  r  r6   r   r   <module>r     sU   5 &  '    	     % > ' 6 ; / A # * ( 4 2 8 < < * 
WW\\"''//(3Y?

  ')*+	A2!F !	FP<~E9V 0434:z	& 	 '+,1,-13$(yP~79 H#L.$-<<4V Vr$(r   