o
    j"Y                     @   s   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	m
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mZmZmZ d dlmZmZmZmZmZmZmZmZ d dlmZ edddZG dd dZd	S )    )SMBDCEFactory)SMBConnection)DCERPC5Connection)samr)RPC_SIDDACL_SECURITY_INFORMATION)NTStatus)DCERPCException)loggerN)
pprint_exc)RPC_C_AUTHN_LEVEL_PKT_INTEGRITYRPC_C_AUTHN_LEVEL_PKT_PRIVACYr	   RPC_C_AUTHN_GSS_NEGOTIATE)RPC_C_AUTHN_LEVEL_NONERPC_C_AUTHN_LEVEL_CONNECTRPC_C_AUTHN_LEVEL_CALLRPC_C_AUTHN_LEVEL_PKTr   r   r	   r   )asynccontextmanagerTFc                 C  sR   t j| |||dI d H \}}|r|z|V  W | I d H  d S | I d H  w )N
auth_levelopenperform_dummy)SAMRRPCfrom_smbconnectionclose)
connectionr   r   r   instanceerr r   /root/aizidognhua/tmp/workspace/projects/ec89d86c-575f-41c9-af57-ac45cbdbf775/venv/lib/python3.10/site-packages/aiosmb/dcerpc/v5/interfaces/samrmgr.pysamrrpc_from_smb   s   "r    c                   @   s  e Zd Zdd Zdd Zdd Zed`d
ededefddZ	ed`d
e
dedefddZdd Zdd Zdd ZejjfddZdd Zdd Zdd Zejfdd Zejfd!efd"d#Zd$d% Zd&d' Zd(d) Zd*d+ Zejfd,d-Zejj fd.d/Z!d0d1 Z"d2d3 Z#ejfd4d5Z$d6d7 Z%ejfd8d9Z&d:d; Z'e(fd<d=Z)ejfd>d?Z*ej+fd@edAefdBdCZ,dDdE Z-dFdG Z.dHdI Z/dJdK Z0dLdM Z1dNdO Z2dPdQ Z3dadRdSZ4dadTdUZ5dVdW Z6e7ejfdXdYZ8e7ejejfdZd[Z9e7ejejfd\d]Z:e7ejejfd^d_Z;dS )br   c                 C   s<   d| _ tj| _d | _d | _i | _i | _i | _i | _	i | _
d S )Nz\samr)service_pipenamer   MSRPC_UUID_SAMRservice_uuiddcehandle
domain_idsdomain_handlesuser_handlesalias_handlesgroup_handlesselfr   r   r   __init__"   s   
zSAMRRPC.__init__c                    s   | S Nr   r+   r   r   r   
__aenter__/   s   zSAMRRPC.__aenter__c                    s   |   I d H  dS NTN)r   )r,   exc_typeexc	tracebackr   r   r   	__aexit__2   s   zSAMRRPC.__aexit__NTFr   r   r   c              
      s   zPt  }| |_|j| |d u r|jt |j I d H \}}|d ur(||j|jI d H \}}|d ur:||du rM| I d H \}}|d urM||d fW S  tyg } z
d|fW  Y d }~S d }~ww NTF)	r   r$   set_auth_levelr   connectbindr#   r   	Exception)r   r   r   r   service_r   er   r   r   from_rpcconnection6   s,   
zSAMRRPC.from_rpcconnectionc              
      s   z)|du rt }t| t jd}tj| |||dI dH \}}|dur&||dfW S  ty@ } z
d|fW  Y d}~S d}~ww )z
		Creates the connection to the service using an established SMBConnection.
		This connection will use the given SMBConnection as transport layer.
		N)filenamer   )r   r   r   r!   r>   get_dce_rpcr:   )r   r   r   r   rpctransportr;   r   r=   r   r   r   r   Q   s    
zSAMRRPC.from_smbconnectionc                    sV   | j s|  I d H \}}|d ur|t| j I d H \}}|d ur$||d | _dS )NServerHandler1   )r$   r8   r   hSamrConnectr%   )r,   r<   r   ansr   r   r   r   d   s   
zSAMRRPC.openc                    s   | j rz| jD ]}z
| |I d H  W q   Y q| jD ]}z
| |I d H  W q   Y q| jD ]}z
| |I d H  W q1   Y q1| jD ]}z
| |I d H  W qF   Y qFz| | jI d H  W n   Y z| j  I d H  W dS    Y dS dS r0   )r$   r'   close_handler(   r*   r)   r%   
disconnect)r,   dhandleuhandler   r   r   r   r   s@   



zSAMRRPC.closec                    s2   t | j|I d H \}}|d urd |fS |d fS r.   )r   hSamrCloseHandler$   )r,   r%   respr   r   r   r   rE      s
   zSAMRRPC.close_handlec                    s6   t j| j||dI d H \}}|d urd |fS |d fS )N)domainInformationClass)r   hSamrQueryInformationDomainr$   )r,   domain_handlerK   rJ   r   r   r   r   get_info   s
   zSAMRRPC.get_infoc                    s^   t | j||gI d H \}}|d urd |fS g }|d d D ]	}||d  q|d |fS )NRelativeIdsElementDatar   )r   hSamrLookupNamesInDomainr$   append)r,   rM   usernamerJ   r   ridsridrawr   r   r   get_user_rid_by_name   s   zSAMRRPC.get_user_rid_by_namec              
   C  s   zOt j}d}|t jkrNtj| j| j|dI d H \}}|d ur,|jt jjkr(|| }|d d D ]	}|d d fV  q2|d }t |d }|t jksW d S W d S  t	yh } zd |fV  W Y d }~d S d }~ww )Nr   enumerationContextBufferNameEnumerationContext	ErrorCode)
r   MORE_ENTRIESr    hSamrEnumerateDomainsInSamServerr$   r%   
error_codevalue
get_packetr:   )r,   statusrY   rJ   r   domainr=   r   r   r   list_domains   s&   
zSAMRRPC.list_domainsc              
      s   z(t | j| j|I d H \}}|d ur||d | j|d  < |d  d fW S  ty? } z
d |fW  Y d }~S d }~ww )NDomainId)r   hSamrLookupDomainInSamServerr$   r%   r&   formatCanonicalr:   )r,   domain_namerJ   r   r=   r   r   r   get_domain_sid   s   zSAMRRPC.get_domain_sidc              
      s~   z't j| j| j| j| |dI d H \}}|d ur||| j|d < |d d fW S  ty> } z
d |fW  Y d }~S d }~ww )NdomainIddesiredAccessDomainHandle)r   hSamrOpenDomainr$   r%   r&   r'   r:   )r,   
domain_sidaccess_levelrJ   r   r=   r   r   r   open_domain   s   &zSAMRRPC.open_domainri   c           	   
      s
  zmd }|   2 z3 d H W \}}|d ur|||kr|} nq6 |d u r*td| t| j| j|I d H \}}|d ur>||d }tj| j| j||dI d H \}}|d urX|| | j|d < || j	| < |d d fW S  ty } z
d |fW  Y d }~S d }~ww )NzDomain not found! %srf   rk   rn   )
re   r:   r   rg   r$   r%   ro   rh   r'   r&   )	r,   ri   rq   	foundname
domainnamer   rJ   	domain_idr=   r   r   r   open_domain_by_name   s6    zSAMRRPC.open_domain_by_namec           
   
   C  s   z^t j}tj}d}|tjkr]t j| j|||dI d H \}}|d ur/|jtjjkr+|| }|d d D ]}d| j	| |d f }|d |d fV  q5|d }t|d }|tjksW d S W d S  t
yx }	 zd d |	fV  W Y d }	~	d S d }	~	ww 	Nr   rX   rZ   %s-%s
RelativeIdr[   r\   r]   )r   USER_NORMAL_ACCOUNTr   r^   hSamrEnumerateUsersInDomainr$   r`   ra   rb   r'   r:   )
r,   rM   	user_typerc   rY   rJ   r   useruser_sidr=   r   r   r   list_domain_users   s*   
zSAMRRPC.list_domain_usersc           	   
   C     zZt j}d}|t jkrYtj| j||dI d H \}}|d ur+|jt jjkr'|| }|d d D ]}d| j| |d f }|d |d fV  q1|d }t |d }|t jksW d S W d S  t	yt } zd d |fV  W Y d }~d S d }~ww rw   )
r   r^   r   hSamrEnumerateGroupsInDomainr$   r`   ra   rb   r'   r:   )	r,   rM   rc   rY   rJ   r   group	group_sidr=   r   r   r   list_domain_groups  (   
zSAMRRPC.list_domain_groupsc              
      s   z.t | j||I d H \}}|d ur |jtjjkr|| }t|d }|tjk}|d fW S  t	yE } z
d |fW  Y d }~S d }~ww )Nr]   )
r   hSamrAddMemberToAliasr$   r`   r   r^   ra   rb   SUCCESSr:   )r,   alias_handlesidrJ   r   rc   resultr=   r   r   r   add_member_to_alias  s   

zSAMRRPC.add_member_to_aliasc           	   
   C  r   rw   )
r   r^   r   r{   r$   r`   ra   rb   r'   r:   )	r,   rM   rc   rY   rJ   r   r}   r~   r=   r   r   r   enumerate_users(  r   zSAMRRPC.enumerate_usersc              
      |   z&t j| j|||dI d H \}}|d ur|| j| | j|d < |d d fW S  ty= } z
d |fW  Y d }~S d }~ww )N)userIdrm   
UserHandle)r   hSamrOpenUserr$   r'   r(   r:   )r,   rM   user_idrq   rJ   r   r=   r   r   r   	open_user;     zSAMRRPC.open_userc              
      sV   zt j| j||dI d H \}}|d fW S  ty* } z
d |fW  Y d }~S d }~ww )N)userInformationClass)r   hSamrQueryInformationUserr$   r:   )r,   user_handler   rJ   r   r=   r   r   r   get_user_infoE  s   
zSAMRRPC.get_user_infoc              
   C  s   z,t | j|I d H \}}|d ur||d d D ]}d| j| |d f d fV  qW d S  tyE } zd |fV  W Y d }~d S d }~ww )NGroupsrx   ry   )r   hSamrGetGroupsForUserr$   r(   r:   )r,   r   rJ   r   r   r=   r   r   r   get_user_group_membershipsM  s   z"SAMRRPC.get_user_group_membershipsc              
   C  s   zQt j}d}|t jkrPtj| j||dI d H \}}|d ur+|jt jjkr'|| }|d d D ]}|d |d d fV  q1|d }t |d }|t jksW d S W d S  tyk } zd d |fV  W Y d }~d S d }~ww )Nr   rX   rZ   r[   ry   r\   r]   )	r   r^   r   hSamrEnumerateAliasesInDomainr$   r`   ra   rb   r:   )r,   rM   rc   rY   rJ   r   aliasr=   r   r   r   list_aliasesX  s&   
zSAMRRPC.list_aliasesc              
      r   )NaliasIdrm   AliasHandle)r   hSamrOpenAliasr$   r'   r)   r:   )r,   rM   alias_idrq   rJ   r   r=   r   r   r   
open_aliask  r   zSAMRRPC.open_aliasc              
   C     z't | j|I d H \}}|d ur||d d D ]}|d  d fV  qW d S  ty@ } zd |fV  W Y d }~d S d }~ww NMembersSids
SidPointer)r   hSamrGetMembersInAliasr$   rh   r:   )r,   r   rJ   r   sidrr=   r   r   r   list_alias_membersu     zSAMRRPC.list_alias_membersc              
      r   )Nr   GroupHandle)r   hSamrOpenGroupr$   r'   r*   r:   )r,   rM   group_idrq   rJ   r   r=   r   r   r   
open_group  r   zSAMRRPC.open_groupc              
   C  r   r   )r   hSamrGetMembersInGroupr$   rh   r:   )r,   group_handlerJ   r   r   r=   r   r   r   list_group_members  r   zSAMRRPC.list_group_membersc              
      s`   zt | j||I d H \}}|d ur||d fW S  ty/ } z
d |fW  Y d }~S d }~ww r.   )r   hSamrQuerySecurityObjectr$   r:   )r,   r%   securityInformationrJ   r   r=   r   r   r   get_security_info  s   
zSAMRRPC.get_security_infoc              
      s   zYt | j| jI d H \}}|d u r||d d d d }t | j| j|I d H \}}|d ur3||d }t j| j| j||dI d H \}}|d urM||| j|d < |d d fW S  typ } z
d |fW  Y d }~S d }~ww )NrZ      r[   rf   rk   rn   )r   r_   r$   r%   rg   ro   r'   r:   )r,   rq   rJ   r   ri   ru   r=   r   r   r   openBuiltinDomain  s&    zSAMRRPC.openBuiltinDomainrT   passwordc           
   
      s   zPt j| j|||t jdI d H \}}|d ur||d }zt | j||dI d H \}}|d ur2|W n t jyH }	 zt | j|  d }	~	ww | |I d H W S  tyg }	 z
d |	fW  Y d }	~	S d }	~	ww )N)accountTyperm   r    )	r   hSamrCreateUser2InDomainr$   USER_ALL_ACCESShSamrSetNTInternal1DCERPCSessionErrorhSamrDeleteUserenable_accountr:   )
r,   rG   rT   r   account_typeresr   r   r<   r=   r   r   r   create_user  s*   "zSAMRRPC.create_userc              
      s   zFt | j|t jjI d H \}}|d ur||d d d }t  }t jj|d< |t jA |d d< t | j||I d H \}}|d urE|W dS  t	y] } z
d |fW  Y d }~S d }~ww NrZ   AllUserAccountControltagControlr1   
r   hSamrQueryInformationUser2r$   USER_INFORMATION_CLASSUserAllInformationSAMPR_USER_INFO_BUFFERUserControlInformationUSER_ACCOUNT_DISABLEDhSamrSetInformationUser2r:   )r,   rH   r   r   uacbufferr=   r   r   r   r     "   zSAMRRPC.enable_accountc              
      s   zFt | j|t jjI d H \}}|d ur||d d d }t  }t jj|d< t j|B |d d< t | j||I d H \}}|d urE|W dS  t	y] } z
d |fW  Y d }~S d }~ww r   r   )r,   rH   r   r   r   r   r<   r=   r   r   r   disable_account  r   zSAMRRPC.disable_accountc                    s   t | j|I d H S r.   )r   r   r$   )r,   rH   r   r   r   delete_account  s   zSAMRRPC.delete_accountc                       t | j|||I d H S r.   )r   hSamrChangePasswordUserr$   )r,   rH   oldpwnewpwr   r   r   change_account_password     zSAMRRPC.change_account_passwordc                    r   r.   )r   hSamrChangePasswordUser_hashr$   )r,   rH   oldpwNTr   r   r   r   change_account_password_nt  r   z"SAMRRPC.change_account_password_ntc                       t | j||I d H S r.   )r   hSamrSetPasswordInternal4Newr$   r,   rH   r   r   r   r   change_account_password_4     z!SAMRRPC.change_account_password_4c                    r   r.   )r   r   r$   r   r   r   r    change_account_password_internal  r   z(SAMRRPC.change_account_password_internalc              
   C    d }|d u r	i }zz|d u r+i }|  |2 z3 d H W \}}}|d ur%|d||< q6 | ||I d H \}}|d ur<|| |2 zI3 d H W \}	}|d urO|tt|	dd }
|
|v r|
|vrd||
< | j||
||d2 z3 d H W \}	}|d ur} n|	d fV  qo6 qA|	d fV  qA6 W n ty } zt	  d |fV  W Y d }~nd }~ww W |d ur| 
|I d H  | j|= d S d S |d ur| 
|I d H  | j|= w )Nr   -)	all_grids
grids_seen)r   r   r   intstrsplitunroll_groupr:   r4   	print_excrE   r*   )r,   rG   gridr   r   ghandlegnametgr   r   ridr=   r   r   r   r     T   
"
zSAMRRPC.unroll_groupc              
   C  r   )Nr   r   r   )	all_arids
arids_seen)r   r   r   r   r   r   unroll_aliasr:   r4   r   rE   r)   )r,   rG   r   r   r   r   r   r   r   r   r   r=   r   r   r   r     r   zSAMRRPC.unroll_aliasc              
   C  s  zFi }i }ddddd}|  |4 I d H "}| |2 z3 d H W \}}}|d ur-|d||< q6 | |2 z3 d H W \}	}
}|d urG|d||
< q86 | |d4 I d H )}| |2 z3 d H W \}}t|dd }d|t|< q]6 W d   I d H  n1 I d H sw   Y  i }i }|D ]}||v rd||< q||v rd||< q|D ])}
| j||
|d2 z3 d H W \}}|d ur ntt|dd }d||< q6 q|D ])}
| j	||
|d2 z3 d H W \}}|d ur ntt|dd }d||< q6 q| 
|2 z*3 d H W \}}}|d ur|tt|dd }||v r$q|||d fV  q6 W d   I d H  W d S 1 I d H sAw   Y  W d S  tyg } zt  d d d |fV  W Y d }~d S d }~ww )Nr   )i  i  i  i  i<  r   r   )r   )r   )get_domain_handler   r   get_alias_handler   r   r   r   r   r   r   r:   r4   r   )r,   ri   
alias_rids
group_ridsdenied_ridsrG   anamearidr   r   r   ahandler   r   denied_groupsdenied_aliasesdriduridnamer=   r   r   r   list_users_allowed_to_replicateE  s   

(
 
 


41z'SAMRRPC.list_users_allowed_to_replicatec              	   C  s   zB|d u s|dks|dkr| j |dI d H \}}n| j||dI d H \}}|d ur,||V  W |d urB| |I d H  | j|= d S d S |d urT| |I d H  | j|= w )Nr   Builtinrq   )r   rv   rE   r'   )r,   ri   rq   rG   r   r   r   r   r     s   
zSAMRRPC.get_domain_handlec              	   C  s   d }zL| j ||d4 I d H !}| j|||dI d H \}}|d ur#||V  W d   I d H  n1 I d H s6w   Y  W |d urN| |I d H  | j|= d S d S |d ur`| |I d H  | j|= w Nr  )r   r   rE   r)   )r,   ri   	alias_riddomain_access_levelalias_access_levelr   rG   r   r   r   r   r     s    (
zSAMRRPC.get_alias_handlec              	   C  s   d }zM| j ||d4 I d H "}| j| |||dI d H \}}|d ur$||V  W d   I d H  n1 I d H s7w   Y  W |d urO| |I d H  | j|= d S d S |d ura| |I d H  | j|= w r	  )r   r   rE   r*   )r,   ri   	group_ridr  group_access_levelr   rG   r   r   r   r   get_group_handle  s    (
zSAMRRPC.get_group_handlec           	   	   C  s   d }z]| j ||d4 I d H 2}| ||I d H \}}|d ur!|| j|||dI d H \}}|d ur4||V  W d   I d H  n1 I d H sGw   Y  W |d ur_| |I d H  | j|= d S d S |d urq| |I d H  | j|= w r	  )r   rW   r   rE   r(   )	r,   ri   	user_namer  user_access_levelrH   rG   user_ridr   r   r   r   get_user_handle  s&   (
zSAMRRPC.get_user_handler6   )NN)<__name__
__module____qualname__r-   r/   r5   staticmethodr   boolr>   r   r   r   r   rE   r   DOMAIN_INFORMATION_CLASSDomainGeneralInformation2rN   rW   re   rj   MAXIMUM_ALLOWEDrr   r   rv   r   r   r   r   r   r   UserGeneralInformationr   r   r   r   r   r   r   r   r   r   rz   r   r   r   r   r   r   r   r   r   r   r  r   r   r   r  r  r   r   r   r   r   !   sd    &




	

')@r   r6   )0aiosmb.dcerpc.v5.common.connection.smbdcefactoryr   aiosmb.connectionr   aiosmb.dcerpc.v5.connectionr   aiosmb.dcerpc.v5r   aiosmb.dcerpc.v5.dtypesr   r   aiosmb.wintypes.ntstatusr   aiosmb.dcerpc.v5.rpcrtr	   aiosmbr
   r4   aiosmb.commons.utils.extbr   r   r   r   r   r   r   r   
contextlibr   r    r   r   r   r   r   <module>   s     (
