
    h                         d Z ddlmZ ddlmZ ddlmZ ddlZddlmZ ddlm	Z	 ddlm
Z
 dd	lmZ dd
lmZ ddlmZ ddlmZ ddlZddddddddddddZdddddddZd d!d"d#Z G d$ d%      Zy)&z)Cloud Workstations configs API utilities.    )absolute_import)division)unicode_literalsN)waiter)GetClientInstance)GetMessagesModule)VERSION_MAP)base)log)	resourceszJ{location}-docker.pkg.dev/cloud-workstations-images/predefined/base:latestzK{location}-docker.pkg.dev/cloud-workstations-images/predefined/clion:latestzN{location}-docker.pkg.dev/cloud-workstations-images/predefined/code-oss:latestzS{location}-docker.pkg.dev/cloud-workstations-images/predefined/code-oss-cuda:latestzL{location}-docker.pkg.dev/cloud-workstations-images/predefined/goland:latestzW{location}-docker.pkg.dev/cloud-workstations-images/predefined/intellij-ultimate:latestzN{location}-docker.pkg.dev/cloud-workstations-images/predefined/phpstorm:latestzM{location}-docker.pkg.dev/cloud-workstations-images/predefined/pycharm:latestzK{location}-docker.pkg.dev/cloud-workstations-images/predefined/rider:latestzN{location}-docker.pkg.dev/cloud-workstations-images/predefined/rubymine:latestzN{location}-docker.pkg.dev/cloud-workstations-images/predefined/webstorm:latest)z
base-imageclioncodeosszcodeoss-cudagolandintellijphpstormpycharmriderrubyminewebstormidmachineTypepoolSizebootDiskSizeGbenableNestedVirtualizationreservationAffinity)r   zmachine-typez	pool-sizezboot-disk-sizezenable-nested-virtualizationreservation-affinitykeyconsumeReservationTypevalues)r   consume-reservation-typer   c                   N    e Zd ZdZej
                  j                  fdZd Zd Z	y)Configsz4The Configs set of Cloud Workstations API functions.c                     t        j                  |      | _        t        |      | _        t        |      | _        | j                  j                  | _        y )N)	r	   getapi_versionr   clientr   messages9projects_locations_workstationClusters_workstationConfigs_service)selfrelease_tracks     2lib/googlecloudsdk/api_lib/workstations/configs.py__init__zConfigs.__init__X   s?    "}5D#M2DK%m4DMMM 	M    c           
         |j                   j                  j                         j                         }|j	                  d      d   }t        j                  d|      j                  d      }|j	                  d      d   }| j                  j                         }||_
        dj                  |j                        |_        dj                  |j                        |_        |j                   r| j                  j                  j#                  t%        t'        j(                  |j                               D cg c]6  \  }}| j                  j                  j"                  j+                  ||      8 c}}      |_        |j,                  |_        |j0                  |_        | j                  j5                         |_        | j                  j9                         |j6                  _        |j<                  |j6                  j:                  _        |j@                  |j6                  j:                  _!        |jD                  r%|jD                  |j6                  j:                  _#        |jH                  r%|jH                  |j6                  j:                  _%        |jL                  |j6                  j:                  _'        |jP                  |j6                  j:                  _)        | j                  jU                  |jV                  |jX                  |jZ                  	      |j6                  j:                  _.        | j                  j_                  |j`                  
      |j6                  j:                  _1        |jd                  |j6                  j:                  _3        |jh                  |j6                  j:                  _5        |jm                  d      r&|jn                  |j6                  j:                  _8        n&|jr                   |j6                  j:                  _8        |jt                  rY|jv                  rM| j                  jy                  |jt                  |jv                        g}	|	|j6                  j:                  _=        | j|                  t        j                  t        j                  j                        k7  rY| j                  j                         |_E        |j                  rd|j                  _G        |j                  rd|j                  _I        | j|                  t        j                  t        j                  j                        k7  r|j                  r}|j                  D ]m  }
| j                  j                         }|
j                         D ]  \  }}|dk(  s|dk(  rD| j                  jy                  |
j                  dd      |
j                  dd            g|_=        U|dk(  r| j                  j                         }|
j                  di       j                         D ]O  \  }}|dk(  r%| j                  j                  j                  |      }t        |t        j                  |      |       Q ||_Q        t        |t        j                  |      |        |j6                  j:                  j                  j                  |       p | j|                  t        j                  t        j                  j                        k7  r|j                  r| j                  j                         }|j                  j                         D ]O  \  }}|dk(  r%| j                  j                  j                  |      }t        |t        j                  |      |       Q ||j6                  j:                  _Q        |j                  rk|j                  D ]\  }| j                  j                         }|j                         D ]  \  }}t        |||        |j                  j                  |       ^ |j                  s| j                  j                         }d|_[        |j                  dk(  r+| j                  j                  j                  j                  }n*| j                  j                  j                  j                  }| j                  j                  |j                  rdn|j                  |j                  rdnd|j                  ||j                        |_d        |j                  j                  |       |j                  r|j                  D ]  }| j                  j                         }|j                  d      |_[        | j                  j                  |j                  d      |j                  d      |j                  d      |j                  d            |_d        |j                  j                  |        | j                  j                         |_k        |j                  r|j                  |j                  _m        nJ|j                  r>t        j                  |j                        j                  |      |j                  _m        |j                  r|j                  |j                  _q        |j                  r|j                  |j                  _s        |j                  r| j                  j                  j                         }|j                  j                         D ]O  \  }}|j                  j                  | j                  j                  j                  j+                  ||             Q ||j                  _w        |j                  |j                  _y        |j                  |j                  _{        | j                  j                         }|j                  r|j                  |_~        |j                  r|j                  |_        ||_        |j                  r|j                  |_        |j                  r|j                  |_        |j                  r|j                  |_        |j                  r| j                  j8                  j                         }|j                  j                         D ]P  \  }}|j                  j                  | j                  j8                  j                  j+                  ||             R ||j6                  j:                  _        | j                  j                  |||      }| j                  j                  |      }t        j                  j!                  dj                  |             |j"                  r=t        j                  j!                  d j                  |j                               |S t$        j&                  j)                  |j                  d!| j|                  "      }t+        j,                  | j                  | j.                  j0                        }t+        j2                  ||d#j                  |j                              }t        j                  j!                  d$j                  |             |S c c}}w )%zCreate a new workstation configuration.

    Args:
      args: argparse.Namespace, The arguments that this command was invoked
        with.

    Returns:
      Workstation configuration that was created.
    /workstationConfigs/r   /locations/(?P<location>[^/]+)/location   {}sr   valueadditionalProperties)enableSecureBoot
enableVtpmenableIntegrityMonitoringenableConfidentialComputedisable_ssh_to_vmtypecountTaccelerator-typeaccelerator-count r   r    z/homeretainext4)sizeGbfsTypediskTypereclaimPolicysourceSnapshotz
mount-pathz	disk-typezsource-snapshotzsource-imagez	read-only)rI   rK   sourceImagereadOnlyr2   )parentworkstationConfigIdworkstationConfigzCreate request issued for: [{}] Check operation [{}] for status.*workstations.projects.locations.operations
collectionr%   &Waiting for operation [{}] to completezCreated configuration [{}].)CONCEPTSconfigParseRelativeNamesplitresearchgroupr'   WorkstationConfignameformatidle_timeoutidleTimeoutrunning_timeoutrunningTimeoutlabelsLabelsValuesortedsix	iteritemsAdditionalPropertydisable_tcp_connectionsdisableTcpConnectionsmax_usable_workstations_countmaxUsableWorkstationsHosthostGceInstancegceInstancemachine_typer   service_accountserviceAccountservice_account_scopesserviceAccountScopesnetwork_tagstags	pool_sizer   disable_public_ip_addressesdisablePublicIpAddressesGceShieldedInstanceConfigshielded_secure_bootshielded_vtpmshielded_integrity_monitoringshieldedInstanceConfigGceConfidentialInstanceConfigenable_confidential_computeconfidentialInstanceConfigenable_nested_virtualizationr   boot_disk_sizer   IsSpecifiedr>   
disableSshenable_ssh_to_vmaccelerator_typeaccelerator_countAcceleratoracceleratorsr%   r	   r$   r
   ReleaseTrackGAHttpOptionshttpOptions-allow_unauthenticated_cors_preflight_requests+allowedUnauthenticatedCorsPreflightRequestsdisable_localhost_replacementdisableLocalhostReplacementboost_configBoostConfigitemsReservationAffinity%ConsumeReservationTypeValueValuesEnumsetattrRESERVATION_AFFINITY_MAPr   BOOST_CONFIG_MAPboostConfigsappendreservation_affinityallowed_ports	PortRangeallowedPortsno_persistent_storagePersistentDirectory	mountPathpd_reclaim_policyGceRegionalPersistentDiskReclaimPolicyValueValuesEnumRETAINDELETEpd_source_snapshotpd_disk_sizepd_disk_typegcePdpersistentDirectoriesephemeral_directoryEphemeralDirectoryGcePersistentDiskephemeralDirectories	Container	containercontainer_custom_imageimagecontainer_predefined_imageIMAGE_URL_MAPcontainer_commandcommandcontainer_argsargscontainer_envEnvValuer8   envcontainer_working_dir
workingDircontainer_run_as_user	runAsUserCustomerEncryptionKeykms_keykmsKeykms_key_service_accountkmsKeyServiceAccountencryptionKeyenable_audit_agentenableAuditAgent&grant_workstation_admin_role_on_create!grantWorkstationAdminRoleOnCreatereplica_zonesreplicaZonesvm_tagsVmTagsValuevmTagsOWorkstationsProjectsLocationsWorkstationClustersWorkstationConfigsCreateRequestr)   Creater   statusPrintasync_r   REGISTRYParseRelativeNamer   CloudOperationPollerr&   projects_locations_operationsWaitFor)r*   r   config_namerO   r2   	config_idrX   r   r6   r   r   desired_boost_configdesired_reservation_affinity
port_rangedesired_port_rangepdreclaim_policy	directoryenv_valencryption_keytags_val
create_reqop_refop_resourcepollerresults                             r,   r   zConfigs.Create`   s    --&&,,.;;=K56q9Fyy;[IOOH !!"89!<I]],,.FFKd&7&78F!LL)=)=>F{{mm55AA
 !'s}}T[['A B	  !C*#u mm--99LL M  !C	  B fm $(#?#?F #'#E#EF  --$$&FK"mm779FKK*.*;*;FKK'-1-A-AFKK*""595P5Pfkk2%)%6%6fkk"'+~~FKK$(( KK4 	//!66))&*&H&H 	0 	
 KK2 	33&*&F&F 	4 	
 KK6 	)) KK6 .2-@-@FKK*+,+/+A+Afkk(/3/D/D+Dfkk(!7!7
--
#
#((** $ l .:fkk*;??4+<+<+?+?@@==446f		;	;IMF		+	+9=6 	KOOD,=,=,@,@AA++,#}}88:&,,.JC&&#1D*D))%))*<bA&**+>B * 1 - ,,113 ) +..&eg
U 2299__ .*..s3 - !4 (*:*>*>s*CUK9 /: 	,,334HI? ,D 	KOOD,=,=,@,@AA%%%)]]%F%F%H"11779*#u,,--33YY% 	($((-	
 : ' kk1 ***!]]446$**,JC
$c5
1 -""#56	 + %%==,,.bbl			8	+MM33PPWW 	
 MM33PPWW 	 88--43D3D..F$$&00 9 bh ""))"- //)]]--/ }}\2==22]];/$==):;!n5]];/	 3 
 	##**2. 0 }}..0F""#::f		(	(,00

)
) !  !%!7!7f"11f''002g**002*#u$$++MM##,,??u @ 	
 3 %f"&"<"<F!%!;!;F ]]88:N||"lln##,0,H,Hn))F $ 7 7f22

5
5 .  ..f||**668h**,*#u%%,,MM%%11DDu E 	
 - (0fkk$nn9 o J ]]!!*-FJJ6==iHI{{	jj9@@MNm$$66?$$ 7 K
 ((t{{@@F ^^077DF
 JJ299)DEM[ s   7;{
c           
      $   |j                   j                  j                         j                         }t	        j
                  d|      j                  d      }|j                  d      d   }| j                  j                         }||_
        | j                  j                  |      }| j                  j                  |      }g }|j                  d      r1dj                  |j                         |_        |j%                  d       |j                  d      r1dj                  |j&                        |_        |j%                  d       |j                  d	      r| j                  j                  j+                  t-        t/        j0                  |j2                              D 	
cg c]6  \  }	}
| j                  j                  j*                  j5                  |	|

      8 c}
}	      |_        |j%                  d	       |j                  d      r"|j6                  |_        |j%                  d       | j:                  t=        j>                  t@        jB                  jD                        k7  r| j                  jG                         |_$        |jJ                  r"d|jH                  _&        |j%                  d       |jN                  r"d|jH                  _&        |j%                  d       |jP                  r"d|jH                  _)        |j%                  d       |jT                  r"d|jH                  _)        |j%                  d       | j                  jW                         |_,        | j                  j[                         |jX                  _.        |j                  d      r6|j^                  |jX                  j\                  _0        |j%                  d       | j:                  t=        j>                  t@        jB                  jD                        k7  r|j                  d      r| j                  jc                         }|jd                  jg                         D ]O  \  }	}
|	dk(  r%| j                  jb                  ji                  |
      }
tk        |tl        j?                  |	      |
       Q ||jX                  j\                  _7        |j%                  d       |j                  d      r6|jp                  |jX                  j\                  _9        |j%                  d       |j                  d      r6|jt                  |jX                  j\                  _;        |j%                  d       |j                  d      r6|jx                  |jX                  j\                  _=        |j%                  d       |j                  d      r6|j|                  |jX                  j\                  _?        |j%                  d       |j                  d      r6|j                  |jX                  j\                  _A        |j%                  d        |j                  d!      r6|j                  |jX                  j\                  _C        |j%                  d"       |j                  d#      r6|j                  |jX                  j\                  _F        |j%                  d$       |j                  d%      r7|j                   |jX                  j\                  _F        |j%                  d$       |j                  d&      rP| j                  j                  |j                  '      |jX                  j\                  _J        |j%                  d(       |j                  d)      r"|j                  |_L        |j%                  d)       |j                  d*      r"|j                  |_N        |j%                  d*       |j                  d+      r"|j                  |_P        |j%                  d+       |j                  d,      r#|j                   |_P        |j%                  d+       |j                  d-      r6|j                  |jX                  j\                  _S        |j%                  d.       | j                  j                         }|j                  d/      r"|j                  |_V        |j%                  d0       |j                  d1      r"|j                  |_X        |j%                  d2       |j                  d3      r"|j                  |_Z        |j%                  d4       ||jX                  j\                  _[        |j                  d5      s|j                  d6      r^| j                  j                  |j                  |j                  7      g}||jX                  j\                  __        |j%                  d8       | j:                  t=        j>                  t@        jB                  jD                        k7  rw|j                  d9      re|j                  D ]D  }| j                  j                         }|jg                         D ]  \  }	}
|	d:k(  s|	d;k(  r*| j                  j                  |d:   |d;   7      g|__        :|	d<k(  r| j                  jc                         }|d<   jg                         D ]O  \  }	}
|	dk(  r%| j                  jb                  ji                  |
      }
tk        |tl        j?                  |	      |
       Q ||_7        tk        |t        j?                  |	      |
        |jX                  j\                  j                  j%                  |       G |j%                  d=       |j                  d>      rg |_d        |j                  D ]\  }| j                  j                         }|jg                         D ]  \  }	}
tk        ||	|
        |j                  j%                  |       ^ |j%                  d>       | j                  j                         |_h        |j                  d?      r-|j                  |j                  _j        |j%                  d@       n`|j                  dA      rOt        j?                  |j                        j                  |B      |j                  _j        |j%                  d@       |j                  dC      r,|j                  |j                  _n        |j%                  dD       |j                  dE      r,|j                  |j                  _p        |j%                  dF       |j                  dG      r| j                  j                  j                         }|j                  jg                         D ]O  \  }	}
|j                  j%                  | j                  j                  j                  j5                  |	|

             Q ||j                  _t        |j%                  dH       |j                  dI      r,|j                  |j                  _v        |j%                  dJ       |j                  dK      r,|j                  |j                  _x        |j%                  dL       |j                  dM      s|j                  dN      r|j                  |_y        |j                  s | j                  j                         g|_y        | j                  j                  |j                  |j                  O      |j                  dP   _~        |j%                  dQ       n|j                  dR      r|j                  |_y        |j                  s | j                  j                         g|_y        | j                  j                  dPdS|j                  T      |j                  dP   _~        |j%                  dQ       |j                  dU      r| j                  jZ                  j                         }|j                  jg                         D ]P  \  }	}
|j                  j%                  | j                  jZ                  j                   j5                  |	|

             R ||jX                  j\                  _        |j%                  dV       |st        j                  dW       yX| j                  j                  ||dYj                  |      Z      }| j                  j                  |      }t        j                  j                  d[j                  |             |j                  r=t        j                  j                  d\j                  |j                               |S t        j                  j                  |j                  d]| j:                  ^      }t        j                  | j                  | j                   j"                        }t        j$                  ||d_j                  |j                              }t        j                  j                  d`j                  |             |S c c}
}	w )azUpdates an existing workstation configuration.

    Args:
      args: argparse.Namespace, The arguments that this command was invoked
        with.

    Returns:
      Workstation configuration that was updated.
    r1   r2   r0   r3   )r`   rb   r4   rd   rf   r5   r7   rn   max_usable_workstationsTz<http_options.allowed_unauthenticated_cors_preflight_requestsFz*http_options.disable_localhost_replacementrt   zhost.gce_instance.machine_typer   r    z&host.gce_instance.reservation_affinityru   z!host.gce_instance.service_accountrw   z(host.gce_instance.service_account_scopesry   zhost.gce_instance.tagsr{   zhost.gce_instance.pool_sizer|   z-host.gce_instance.disable_public_ip_addressesr   z#host.gce_instance.boot_disk_size_gbr>   zhost.gce_instance.disable_sshr   r   r<   zJhost.gce_instance.confidential_instance_config.enable_confidential_computer   r   rl   enable_tcp_connectionsr   z.host.gce_instance.enable_nested_virtualizationr   z=host.gce_instance.shielded_instance_config.enable_secure_bootr   z6host.gce_instance.shielded_instance_config.enable_vtpmr   zFhost.gce_instance.shielded_instance_config.enable_integrity_monitoringr   r   r?   zhost.gce_instance.acceleratorsr   rB   rC   r   zhost.gce_instance.boost_configsr   r   zcontainer.imager   rN   r   zcontainer.commandr   zcontainer.argsr   zcontainer.envr   zcontainer.working_dirr   zcontainer.run_as_userr   r   )rG   rI   r   persistent_directoriesr   rD   )rG   rH   rK   r   zhost.gce_instance.vm_tagszNo fields were specified.N,)r`   rQ   
updateMaskzUpdate request issued for: [{}]rR   rS   rT   rV   zUpdated configuration [{}].)rW   rX   rY   rZ   r\   r]   r^   r[   r'   r_   r`   LWorkstationsProjectsLocationsWorkstationClustersWorkstationConfigsGetRequestr)   Getr   ra   rb   rc   r   rd   re   rg   rh   ri   rj   rf   rk   rn   ro   r%   r	   r$   r
   r   r   r   r   r   r   0disallow_unauthenticated_cors_preflight_requestsenable_localhost_replacementr   r   rp   rq   rr   rs   rt   r   r   r   r   r   r   r   r   ru   rv   rw   rx   ry   rz   r{   r   r|   r}   r   r   IsKnownAndSpecifiedr>   r   r   r   r   r   r   r   r   r   rl   rm   r   r   r   r~   r   r9   r   r:   r   r;   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r8   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   errorNWorkstationsProjectsLocationsWorkstationClustersWorkstationConfigsPatchRequestjoinPatchr   r   r   r   r   r   r   r   r&   r   r   )r*   r   r   r2   r   rX   get_req
old_configupdate_maskr   r6   r   gce_shielded_instance_configr   r   r   r   r   r   r   
update_reqr   r   r   r   s                            r,   UpdatezConfigs.Updatef  s9    --&&,,.;;=Kyy;[IOOH !!"89!<I]],,.FFKmmhh i G ""7+JK' <<(9(9:f()*#ll4+?+?@f*+!mm55AA
 !'s}}T[['A B	  !C*#u mm--99LL M  !C	  B fm "78%)%G%Gf"23;??4+<+<+?+?@@==446f		;	;IMFJ	
 
	>	>INFJ	
 
	*	*9>6GH		+	+9=6GH --$$&FK"mm779FKK',0,=,=fkk)9:;?? 


1
2%)]]%F%F%H"11779*#u,,--33YY% 	($((-	
 : ' kk1 AB)*/3/C/Cfkk,<=01595P5Pfkk2CD'%)%6%6fkk"12$)-fkk&6756

*
* kk6 HI()/3/B/Bfkk,>? 34+/+A+Afkk(89 23/3/D/D+Dfkk(8956
--
5
5(,(H(H 6  kk8
 
V ,- $ 7 7f-.@A

5
5 . AB12%)%A%Af"2301)-)D)D%Df"2367

+
+ kk8 IJ $(==#J#J#L ./6:6O6O"3
I (040B0B"-
B 78

,
, #< 
R
 	% KK2 *+t/?/?0 --
#
#((** $ l .:fkk*9:;?? 


>
*++,#}}88:&,,.JC&&#1D*D))%&89&':; * 1 - ,,+/==+L+L+N(*+ABHHJ
U2299__ .*..s3 K - !4 (*:*>*>s*CUK1 /2 	,,334HI7 ,8 :;(f***!]]446$**,JC
$c5
1 -""#56	 +
 ) }}..0F01#::f*+			6	7,00

)
) !  *++,!%!7!7f,-()"11f)*(''002g**002*#u$$++MM##,,??u @ 	
 3 %f)/0$($>$>f!01/0#'#=#=f 01'4+;+;N+K%/%E%Ef"--(,(I(I(K'L$ --
1
1&&1B1B 2  ""1%+
 12			.	/%/%E%Ef"--(,(I(I(K'L$
--
1
1r$2I2I 2  ""1%+
 12	"**668h**,*#u%%,,MM%%11DDu E 	
 - (0fkk$45	ii+,mm 88K( n J
 ]]  ,FJJ6==iHI{{	jj9@@MNm$$66?$$ 7 K
 ((t{{@@F ^^077DF
 JJ299)DEMo	 s   %;AI
N)
__name__
__module____qualname____doc__r
   r   BETAr-   r   r   r.   r,   r"   r"   U   s&    <#'#4#4#9#9 DLYr.   r"   )r	  
__future__r   r   r   r\   googlecloudsdk.api_lib.utilr   (googlecloudsdk.api_lib.workstations.utilr   r   r	   googlecloudsdk.callioper
   googlecloudsdk.corer   r   ri   r   r   r   r"   r  r.   r,   <module>r     s    0 &  ' 	 . F F @ ( # ) 

 	U 	V 	Y 	^ 	W 	b 	Y 	X 	V 	Y 	YA"J !&$@1   8 j	 j	r.   