
    lmjH2                       d Z ddlmZ ddlZddlZddlZddlZddlZddlZddl	Z	ddl
Z
ddlmZ ddlmZ ddlmZmZ ddlmZmZmZ ddlmZmZ d	d
hZh dZ ej6                  d      Z ej6                  d      Z ej6                  d      Z G d de      Z ddZ!ddZ"ddZ#	 d	 	 	 	 	 	 	 	 	 ddZ$d dZ%d!dZ&	 d	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 d"dZ'ddZ(d#dZ)ef	 	 	 	 	 	 	 	 	 d$dZ*y)%zServer-side connection checks for external integrations.

Credential files are referenced by absolute path. Secret values never cross the
admin HTTP API and are never included in errors or successful responses.
    )annotationsN)Path)Callable)	HTTPErrorURLError)quoteunquoteurlparse)Requesturlopenpostgres
postgresql>   	verify-caverify-fullallowpreferdisablerequirez ^[a-z0-9][a-z0-9-]{1,62}-[0-9]+$z^[a-z][a-z0-9-]{1,30}$z7\.cos\.(?P<region>[a-z][a-z0-9-]{1,30})\.myqcloud\.com$c                  "     e Zd Zdd fdZ xZS )IntegrationAdapterErrorc                N    t         |   |       || _        || _        || _        y N)super__init__codemessagehttp_status)selfr   r   r   	__class__s       3/opt/aitasker-platform-test/integration_adapters.pyr   z IntegrationAdapterError.__init__    s&    !	&    )  )r   strr   r#   r   intreturnNone)__name__
__module____qualname__r   __classcell__)r   s   @r    r   r      s    ' 'r!   r   c                6   t         j                  j                  |d      j                         xs | j                         }|st	        dd| dd      t        |      }|j                         st	        d| dd      	 t        j                  |j                  d	            }t        |t              st	        d| dd      |S # t        $ r}t	        d
| dd      |d }~wt        t        t        j                  f$ r}t	        d| dd      |d }~ww xY w)N CONFIG_MISSINGu	   未配置u   凭据文件  INVALID_CREDENTIAL_REFu!   凭据引用必须是绝对路径utf-8)encodingCREDENTIAL_FILE_NOT_FOUNDu   凭据文件不存在CREDENTIAL_FILE_INVALIDu'   凭据文件不可读取或格式错误u   凭据文件格式错误)osenvirongetstripr   r   is_absolutejsonloads	read_textFileNotFoundErrorOSErrorUnicodeDecodeErrorJSONDecodeError
isinstancedict)request_refenvironment_keylabelrefcredential_pathpayloadexcs          r    _read_json_credentialsrI   '   s/   
**.."
-
3
3
5
L9J9J9LC%&6)E7,8WY\]]3iO&&(%&>5'Ij@kmpqqB**_666HI
 gt$%&?E7JbAcehiiN  r%&AeWLaCbdghnqq')=)=> B%&?E7JqArtwx  B  	BBs$    %C 	DC!!!DDDc                   t        | j                               }|j                  t        vs|j                  st        ddd      |j                  |j                  t        ddd      t        |j                  j                  d            j                         }|st        ddd      |j                         r |j                         |k7  rt        dd	d      	 |j                  xs d
}t        |dd      }|j                  d      }|j                  d      }	t        |j                  d      xs d      }
t!        |t              r"|j                         rt!        |	t              r|	st        ddd      |
t"        vrt        ddd      |j                  |||j                         |	|
ddS # t        $ r}t        ddd      |d }~ww xY w)NINVALID_DATABASE_ENDPOINTu=   数据库地址应为 postgresql://主机:端口/数据库名r.   DATABASE_SECRET_IN_ENDPOINTu-   数据库地址中不得填写账号或密码/u!   数据库地址缺少数据库名DATABASE_IDENTIFIER_MISMATCHu3   数据库标识与地址中的数据库名不一致i8  u   数据库端口无效"AITASKER_DATABASE_CREDENTIALS_FILE
PostgreSQLuserpasswordsslmoder   DATABASE_CREDENTIALS_INVALIDu/   PostgreSQL 凭据文件缺少 user 或 passwordu   PostgreSQL sslmode 配置无效   )hostportdbnamerQ   rR   rS   connect_timeout)r
   r7   schemePOSTGRES_SCHEMEShostnamer   usernamerR   r	   pathlstriprW   
ValueErrorrI   r6   r#   r@   POSTGRES_SSL_MODES)endpoint
identifiercredentials_refparseddatabaserW   rH   credentialsrQ   rR   rS   s              r    _postgres_connection_settingsrh   9   s   hnn&'F}},,FOO%'K
 	

 "foo&A%&CEtvyzzv{{))#./557H%&ACfhkllj..0H<%&DF{  ~A  B  	Bj{{"d ),K
 ??6"Dz*H+//),9	:GdC 

JxQT<U]e%*=
 	

 ((%&DFgilmm

 '  j%&ACZ\_`fiijs   F/ /	G8GGc                    t        |       j                         dv sdv rt        ddd      S dv sdv rt        dd	d
      S d}t        fd|D              rt        ddd      S t        ddd      S )Nzpassword authentication failedzauthentication failedDATABASE_AUTH_FAILEDu2   PostgreSQL 认证失败，请检查账号和密码  timeoutz	timed outDATABASE_TIMEOUTu.   PostgreSQL 连接超时，请检查 SSH 隧道  )zconnection refusedzcould not connectzfailed to connectznetwork is unreachablezno connection could be madezserver closed the connectionc              3  &   K   | ]  }|v  
 y wr    ).0markerdetails     r    	<genexpr>z"_database_error.<locals>.<genexpr>v   s     
>6V
>s   DATABASE_UNREACHABLEu7   PostgreSQL 无法连接，请检查 SSH 隧道和端口DATABASE_TEST_FAILEDu    PostgreSQL 连接或查询失败)r#   lowerr   any)rH   unreachable_markersrs   s     @r    _database_errorrz   h   s    X^^F'615LPV5V&'=?suxyyFkV3&'9;kmpqq 
>*=
>>&'=?xz}~~"#9;]_bccr!   c                @   t        | ||      }|	 dd l}|j                  }d }	  |di |}|j                         5 }|j                  d       |j                         }	d d d        	r|	d   dk7  rt        ddd	      	 |	 |j                          d
ddddS # t        $ r}t        ddd      |d }~ww xY w# 1 sw Y   ZxY w# t        $ r  t        $ r}t        |      |d }~ww xY w# t        $ r Y cw xY w# |!	 |j                          w # t        $ r Y w w xY ww xY w)Nr   DATABASE_DRIVER_MISSINGu4   缺少 PostgreSQL 驱动，请安装 psycopg[binary]r"   zSELECT 1   rv   u   PostgreSQL 查询校验失败rk   Trf   	CONNECTEDu+   PostgreSQL 真实连接通过（SELECT 1）okservicestatusr   rp   )rh   psycopgImportErrorr   connectcursorexecutefetchone	Exceptionrz   close)
rb   rc   rd   	connectorsettingsr   rH   
connectionr   rows
             r    test_postgresql_connectionr   {   sn    -Xz?SH	 OO	J**
  	$FNN:&//#C	$ c!fk)*@Bacfgg " !  " @	 5  	))F 			$ 	$
 #  ,c"+,   !  "  "s   B C "B7C C) 	B4!B//B47C <C C&C!!C&&C8 )	C54C58D<DD	DDDDc                   t        | dd      }|j                  d      }|j                  d      }t        |t              r"|j	                         rt        |t              r|st        ddd      t        t        j                  j                  d      xs |j                  d	      xs d
      j	                         }t        |j	                               }|j	                         r(|j                  dk7  s|j                  st        ddd      |s@|j                  r4t        j                  |j                        }|r|j                  d	      nd
}t        j                  |      st        ddd      |j	                         ||fS )NTENCENT_COS_CREDENTIALS_FILEu   腾讯云 COSsecretId	secretKeyCOS_CREDENTIALS_INVALIDu7   腾讯云 COS 凭据文件缺少 secretId 或 secretKeyr.   TENCENT_COS_REGIONregionr,   httpsINVALID_COS_ENDPOINTu"   COS 访问域名必须使用 HTTPSCOS_REGION_MISSINGu7   COS 凭据文件缺少有效 region，例如 ap-beijing)rI   r6   r@   r#   r7   r   r4   r5   r
   rZ   r\   COS_API_HOST_PATTERNsearchgroupCOS_REGION_PATTERN	fullmatch)rd   rb   rg   	secret_id
secret_keyr   re   matchs           r    _cos_credentialsr      sH   (&K
 
+I-Ji%Y__->jQ[]`Fais%%E
 	
  45X9RXVXY__aFhnn&'F~~V]]g5V__%&<>bdghhfoo$++FOO<*/X&R''/% E
 	

 ??j&00r!   c                   t        t        j                               dz
  }| d|dz    }d}| j                          d| d| d}d| dt        j                  |j                  d            j                          d}	t        j                  |j                  d      |j                  d      t        j                        j                         }
t        j                  |
j                  d      |	j                  d      t        j                        j                         }d	||||d
|d}dj                  d |j                         D              S )N   ;iX  rV   
z

host=zsha1
r0   sha1r,   )zq-sign-algorithmzq-akzq-sign-timez
q-key-timezq-header-listzq-url-param-listzq-signature&c              3  F   K   | ]  \  }}| d t        |d         yw)=z;-_.~safeN)r   )rq   keyvalues      r    rt   z%_cos_authorization.<locals>.<genexpr>   s'     [jc5se1U5w789[s   !)r$   timerw   hashlibr   encode	hexdigesthmacnewjoinitems)methodr^   rV   r   r   startkey_timeheader_listhttp_stringstring_to_signsign_key	signaturevaluess                r    _cos_authorizationr      s$   		r!E%#+'HK\\^$BtfIdV2>KhZr',,{7I7I'7R*S*]*]*_)``bcNxx
))'2HOOG4Lgll[eegH1>3H3H3QSZS_S_`jjlI"$ F 88[FLLN[[[r!   c           
     .   | d| d}dt        |d      z   }t        d| | ||       }	|	j                  d|       |	j                  d	t        | ||||             |6|	j                  d
d       |	j                  dt	        t        |                   |	S )Nz.cos.z.myqcloud.comrM   z/-_.~r   zhttps://)datar   HostAuthorizationzContent-Typezapplication/octet-streamzContent-Length)r   r   
add_headerr   r#   len)
r   bucketr   
object_keyr   r   bodyrV   r^   requests
             r    _build_cos_requestr      s     XU6(-0Dz00Dtf-DHGvt$(:64yZd(ef>+EF+ST^<Nr!   c                4   t        | t              r{| j                  dv rt        ddd      S | j                  dk(  rt        ddd      S | j                  dv s| j                  d	k\  rt        d
dd      S t        dd| j                   d      S t        | t        t
        j                  f      rt        ddd      S t        | t              rDt        | j                  t        t
        j                  f      rt        ddd      S t        ddd      S t        ddd      S )N)i  i  COS_ACCESS_DENIEDu+   COS 认证失败或没有测试对象权限rk   i  COS_RESOURCE_NOT_FOUNDu+   COS Bucket、地域或测试对象不存在)i  i  i  COS_TEMPORARILY_UNAVAILABLEu%   COS 暂时不可用，请稍后重试COS_TEST_FAILEDu   COS 返回 HTTP COS_TIMEOUTu   COS 连接超时rn   COS_UNREACHABLEu   COS 网络连接失败u   COS 连接测试失败)	r@   r   r   r   TimeoutErrorsocketrl   r   reason)rH   s    r    
_cos_errorr      s   #y!88z!*+>@morss88s?*+CErtwxx88z!SXX_*+HJqsvww&'8<LSXXJ:WY\]]#fnn56&}6H#NN#x cjj<"@A*=:LcRR&'8:RTWXX"#46NPSTTr!   c                d     | |d      5 }|j                  d      cd d d        S # 1 sw Y   y xY w)N   )rl   i   )read)openerr   responses      r    _open_cos_requestr      s/    		# (x}}Y'( ( (s   &/c           
     l   | j                         } t        j                  |       st        ddd      t	        ||      \  }}}dt        j                         j                   d}dt        j                         j                   j                  d      }d}	d }
	 t        d	| |||||      }t        ||       d
}	t        d| ||||      }t        ||      }||k7  rt        ddd      	 |	r	 t        d| ||||      }t        ||       |
|
d
ddddS # t        $ r}|}
Y d }~:d }~wt        $ r}t        |      }
Y d }~Vd }~ww xY w# t        $ r}|
t        |      }
Y d }~Zd }~ww xY w# |	rF	 t        d| ||||      }t        ||       w # t        $ r}|
t        |      }
Y d }~w Y d }~w d }~ww xY ww xY w)NINVALID_COS_BUCKETu   COS Bucket 名称格式无效r.   z_aitasker_connection_tests/z.txtzaitasker-cos-connection-test:asciiFPUTTGETCOS_CONTENT_MISMATCHu%   COS 上传后读取内容校验失败rk   DELETEstorager~   u>   腾讯云 COS 真实连接通过（上传、读取、删除）r   )r7   COS_BUCKET_PATTERNr   r   r   uuiduuid4hexr   r   r   r   r   )r   rb   rd   r   r   r   r   r   contentuploadedprimary_errorput_requestget_requestreturnedrH   delete_requests                   r    test_cos_connectionr     s    \\^F''/%&:<[]`aa$4_h$O!Iz6.tzz|/?/?.@EJ-djjl.>.>-?@GGPGH48M4(
IWacjk&+.(
IWab$V[9w)*@Biknoo  4!3HffjZceo!p!&.9
  S	  #  ("3(  4 ($.sOM4	 4!3HffjZceo!p!&.9 4 ($.sOM )4	 s   AD 'E) +E 	ED"E) "E.D>9E) >EE) 	E&E!!E&)F3-F
	F3
	F/F* 
F3*F//F3)rB   r#   rC   r#   rD   r#   r%   dict[str, object])rb   r#   rc   r#   rd   r#   r%   r   )rH   r   r%   r   r   )
rb   r#   rc   r#   rd   r#   r   zCallable[..., object] | Noner%   r   )rd   r#   rb   r#   r%   ztuple[str, str, str])r   r#   r^   r#   rV   r#   r   r#   r   r#   r%   r#   )r   r#   r   r#   r   r#   r   r#   r   r#   r   r#   r   zbytes | Noner%   r   )r   Callable[..., object]r   r   r%   bytes)
r   r#   rb   r#   rd   r#   r   r   r%   r   )+__doc__
__future__r   r   r   r9   r4   rer   r   r   pathlibr   typingr   urllib.errorr   r   urllib.parser   r	   r
   urllib.requestr   r   r[   ra   compiler   r   r   RuntimeErrorr   rI   rh   rz   r   r   r   r   r   r   r   rp   r!   r    <module>r      s|   #    	 	      , 1 1 + - Z RZZ CD RZZ 9: !rzz"\] 'l '$,^d. /3	))) ) ,	)
 )X1>\6   	
    (U$( %,	+++ + "	+
 +r!   