o
    j\                     @   sl  d dl mZ d dlZd dlZd dlZd dlZd dlT d dlmZ d dl	T G dd dej
ZG dd dej
ZG d	d
 d
ZG dd dejZG dd dejZG dd dZG dd dZG dd dejZG dd dZG dd dejZG dd dZG dd dejZG dd dejZG dd  d ZG d!d" d"ejZG d#d$ d$ZG d%d& d&ZG d'd( d(ZdS ))    )ContextN)*)NegotiateDialectsc                   @      e Zd ZdZdZdZdS )NegotiateSecurityModer         N)__name__
__module____qualname__NONESMB2_NEGOTIATE_SIGNING_ENABLEDSMB2_NEGOTIATE_SIGNING_REQUIRED r   r   /root/aizidognhua/tmp/workspace/projects/ec89d86c-575f-41c9-af57-ac45cbdbf775/venv/lib/python3.10/site-packages/aiosmb/protocol/smb2/commands/negotiate.pyr          r   c                   @   (   e Zd ZdZdZdZdZdZdZdZ	dS )	NegotiateCapabilitiesr   r                @   N)
r	   r
   r   DFSLEASING	LARGE_MTUMULTI_CHANNELPERSISTENT_HANDLESDIRECTORY_LEASING
ENCRYPTIONr   r   r   r   r          r   c                   @   s0   e Zd Zdd Zdd Zedd Zdd Zd	S )
NEGOTIATE_REQc                 C   sL   d| _ d | _d | _d| _d | _d | _d| _d | _d | _d| _	g | _
g | _d S )N$   r   )StructureSizeDialectCountSecurityModeReservedCapabilities
ClientGuidClientStartTimeNegotiateContextOffsetNegotiateContextCount	Reserved2NegotiateContextListDialectsselfr   r   r   __init__    s   
zNEGOTIATE_REQ.__init__c                 C   s  | j g kr	td| jjdddd}|t| j jdddd7 }|| jjdddd7 }|| jjdddd7 }|| jjdddd7 }|| j 7 }t	j
| j v rd}| j D ]}||jjdddd7 }qRdt| }|d	d
| d
  7 }dt| | _t| j| _|| jjdddd7 }|| jjdddd7 }|| jjdddd7 }||7 }t| jD ]\}}|| 7 }|| jd kr|d	d
t| d
  7 }q|S || jjd
ddd7 }| j D ]}||jjdddd7 }q|S )Nz!At least one dialect MUST be set!r   littleF	byteordersignedr       d       r   r   )r.   	Exceptionr#   to_byteslenr%   r&   r'   r(   r   SMB311valuer*   r-   r+   r,   	enumerater)   )r0   tdpdialectpadding_startsictxr   r   r   r:   /   s<   


zNEGOTIATE_REQ.to_bytesc              	   C   s  t  }tj| dddd|_|jdksJ tj| dddd|_|jdks)J ttj| dddd|_| d|_t	tj| dddd|_
t| |_|  }| dtj g |_td|jD ]}|jttj| dddd qe|  }| |tj tj|jv rtj| dddd|_tj| dddd|_tj| dddd|_g |_| |jtj td|jD ]%}|jt|  |  }t|d\}}|dkr| |d	 d tj q|S tj| dddd|_| |tj |S )
Nr   r2   Fr3   r"   r   r   r   r   )r!   int
from_bytesreadr#   r$   r   r%   r&   r   r'   GUIDfrom_bufferr(   tellseekioSEEK_CURr.   rangeappendr   SEEK_SETr<   r*   r+   r,   r-   SMB2NegotiateContextdivmodr)   )buffcmdpos_pos_buff_endpad_posqmr   r   r   rI   T   sB   $zNEGOTIATE_REQ.from_bufferc                 C   s   d}|d| j  7 }|d| j 7 }|d| jj 7 }|d| j 7 }|dt| j 7 }|d| j 7 }tj	| j
v rW|d| j 7 }|d	| j 7 }|d
| j 7 }| jD ]}|t|7 }qMn|d| j 7 }| j
D ]	}|d|j 7 }qa|S )Nz==== SMB2 NEGOTIATE REQ ====
StructureSize: %s
zDialectCount:  %s
zSecurityMode:  %s
zReserved:      %s
zCapabilities:  %s
zClientGuid:    %s
zNegotiateContextOffset:    %s
zNegotiateContextCount:    %s
zReserved2:    %s
zClientStartTime:    %s
z	 Dialect: %s
)r#   r$   r%   namer&   reprr'   r(   r   r<   r.   r*   r+   r,   r-   r)   )r0   r?   rD   rA   r   r   r   __repr__~   s$   

zNEGOTIATE_REQ.__repr__N)r	   r
   r   r1   r:   staticmethodrI   r^   r   r   r   r   r!      s    %
)r!   c                   @   r   )	SMB2ContextTyper   r               r   N)
r	   r
   r   PREAUTH_INTEGRITY_CAPABILITIESENCRYPTION_CAPABILITIESCOMPRESSION_CAPABILITIESNETNAME_NEGOTIATE_CONTEXT_IDTRANSPORT_CAPABILITIESRDMA_TRANSFORM_CAPABILITIESSIGNING_CAPABILITIESr   r   r   r   r`      r    r`   c                   @   s   e Zd ZdZdS )SMB2HashAlgorithmr   N)r	   r
   r   SHA_512r   r   r   r   rl      s    rl   c                   @   s    e Zd Zdd Zedd ZdS )rQ   c                 C   s   d S Nr   r/   r   r   r   r1      s   zSMB2NegotiateContext.__init__c                 C   s   |   }ttj| dddd}| | |tjkr t| S |tj	kr*t
| S |tjkr4t| S |tjkr>t| S |tjkrHt| S |tjkrRt| S td| )Nr   r2   Fr3   zUnknown contextype %s)rJ   r`   rE   rF   rG   rK   re    SMB2PreauthIntegrityCapabilitiesrI   rf   SMB2EncryptionCapabilitiesrg   SMB2CompressionCapabilitiesrh   SMB2NetnameNegotiateContextIDri   SMB2TransportCapabilitiesrk   SMB2SigningCapabilitiesr9   )rS   rU   ContextTyper   r   r   rI      s    












z SMB2NegotiateContext.from_bufferN)r	   r
   r   r1   r_   rI   r   r   r   r   rQ      s    rQ   c                   @   sF   e Zd Zdd Zedd ZeejgdfddZdd	 Z	d
d Z
dS )ro   c                 C   s0   t j| _d | _d| _d | _d | _d | _d | _d S Nr   )	r`   re   ru   
DataLengthr&   HashAlgorithmCount
SaltLengthHashAlgorithmsSaltr/   r   r   r   r1         
z)SMB2PreauthIntegrityCapabilities.__init__c              	   C   s   t  }ttj| dddd|_tj| dddd|_tj| dddd|_tj| dddd|_tj| dddd|_	g |_
t|jD ]}|j
ttj| dddd qI| |j	|_|S Nr   r2   Fr3   r   )ro   r`   rE   rF   rG   ru   rw   r&   rx   ry   rz   rN   rO   rl   r{   rS   caprV   r   r   r   rI      s   $z,SMB2PreauthIntegrityCapabilities.from_bufferr6   c                 C   s   t  }| |_||_|S rn   )ro   rz   r{   )
hash_algossaltr   r   r   r   	construct   s   z*SMB2PreauthIntegrityCapabilities.constructc                 C   s   t | j| _t | j| _| jjdddd}|| jjdddd7 }| jD ]}||jjdddd7 }q#|| j7 }t || _| jjjdddd}|| jjdddd7 }|| j	jdddd7 }||7 }|S r}   )
r;   r{   ry   rz   rx   r:   r=   rw   ru   r&   )r0   datahashalgor?   r   r   r   r:      s   


z)SMB2PreauthIntegrityCapabilities.to_bytesc                 C   sL   d}|d| j  7 }|d| j 7 }|d| j 7 }| jD ]	}|d|j 7 }q|S )Nz/==== SMB2 Preauth Integrity Capabilities ====
zHashAlgorithmCount: %s
zSaltLength: %s
z
Salt: %s
zHashAlgo: %s
)rx   ry   r{   rz   r\   )r0   r?   algor   r   r   r^      s   
z)SMB2PreauthIntegrityCapabilities.__repr__N)r	   r
   r   r1   r_   rI   rl   rm   r   r:   r^   r   r   r   r   ro      s    

ro   c                   @   s   e Zd ZdZdZdZdZdS )
SMB2Cipherr   r   ra   r   N)r	   r
   r   AES_128_CCMAES_128_GCMAES_256_CCMAES_256_GCMr   r   r   r   r     s
    r   c                   @   <   e Zd Zdd Zedd Zedd Zdd Zd	d
 ZdS )rp   c                 C   $   t j| _d | _d| _d | _d | _d S rv   )r`   rf   ru   rw   r&   CipherCountCiphersr/   r   r   r   r1   
  
   
z#SMB2EncryptionCapabilities.__init__c              	   C      t  }ttj| dddd|_tj| dddd|_tj| dddd|_tj| dddd|_g |_	t
|jD ]}|j	ttj| dddd q=|S r}   )rp   r`   rE   rF   rG   ru   rw   r&   r   r   rN   rO   r   r~   r   r   r   rI        $z&SMB2EncryptionCapabilities.from_bufferc                 C      t  }| |_t| |_|S rn   )rp   r   r;   r   )supp_enc_listsr   r   r   from_enc_list      
z(SMB2EncryptionCapabilities.from_enc_listc                 C      | j jdddd}| jD ]}||jjdddd7 }qt|| _| jjjdddd}|| jjdddd7 }|| jjdddd7 }||7 }|S r}   )r   r:   r   r=   r;   rw   ru   r&   )r0   r   cipherr?   r   r   r   r:   (     

z#SMB2EncryptionCapabilities.to_bytesc                 C   0   d}|d| j  7 }| jD ]	}|d|j 7 }q|S )Nz(==== SMB2 Encryption Capabilities ====
zCipherCount: %s
zCipher: %s
)r   r   r\   )r0   r?   r   r   r   r   r^   7  
   
z#SMB2EncryptionCapabilities.__repr__N	r	   r
   r   r1   r_   rI   r   r:   r^   r   r   r   r   rp   	      

rp   c                   @   r   )SMB2SigningAlgorithmr   r   r   N)r	   r
   r   HMAC_SHA256AES_CMACAES_GMACr   r   r   r   r   ?  r   r   c                   @   r   )rt   c                 C   r   rv   )r`   rk   ru   rw   r&   SigningAlgorithmCountSigningAlgorithmsr/   r   r   r   r1   F  r   z SMB2SigningCapabilities.__init__c              	   C   r   r}   )rt   r`   rE   rF   rG   ru   rw   r&   r   r   rN   rO   r   r~   r   r   r   rI   N  r   z#SMB2SigningCapabilities.from_bufferc                 C   r   rn   )rt   r   r;   r   )supp_sig_listr   r   r   r   r   \  r   z%SMB2SigningCapabilities.from_enc_listc                 C   r   r}   )r   r:   r   r=   r;   rw   ru   r&   )r0   r   sigr?   r   r   r   r:   d  r   z SMB2SigningCapabilities.to_bytesc                 C   r   )Nz%==== SMB2 Singing Capabilities ====
zSigningAlgorithmCount: %s
z	Sig: %s
)r   r   r\   )r0   r?   r   r   r   r   r^   s  r   z SMB2SigningCapabilities.__repr__Nr   r   r   r   r   rt   E  r   rt   c                   @   s    e Zd ZdZdZdZdZdZdS )SMB2CompressionTyper   r   r   ra   r   N)r	   r
   r   r   LZNT1LZ77LZ77_Huffman
Pattern_V1r   r   r   r   r   {  s    r   c                   @      e Zd ZdZdZdS )SMB2CompressionFlagsr   r   N)r	   r
   r   r   CHAINEDr   r   r   r   r         r   c                   @   s>   e Zd Zdd Zedd ZedddZdd	 Zd
d ZdS )rq   c                 C   s0   t j| _d | _d| _d | _d| _d | _g | _d S rv   )	r`   rg   ru   rw   r&   CompressionAlgorithmCountPaddingFlagsCompressionAlgorithmsr/   r   r   r   r1     r|   z$SMB2CompressionCapabilities.__init__c              	   C   s   t  }ttj| dddd|_tj| dddd|_tj| dddd|_tj| dddd|_tj| dddd|_	t
tj| dddd|_t|jD ]}|jttj| dddd qT|S r}   )rq   r`   rE   rF   rG   ru   rw   r&   r   r   r   r   rN   r   rO   r   r~   r   r   r   rI     s   $z'SMB2CompressionCapabilities.from_bufferFc                 C   s2   t  }tj|_|du rtj|_| |_t| |_|S )NT)rq   r   r   r   r   r   r;   r   )supp_comp_list
is_chainedr   r   r   r   from_comp_list  s   
z*SMB2CompressionCapabilities.from_comp_listc                 C   s   | j jdddd}|| jjdddd7 }|| jjjdddd7 }|| j jdddd7 }| jD ]}||jjdddd7 }q.t|| _| jjjdddd}|| jjdddd7 }|| j	jdddd7 }||7 }|S r}   )
r   r:   r   r   r=   r   r;   rw   ru   r&   )r0   r   compalgor?   r   r   r   r:     s   

z$SMB2CompressionCapabilities.to_bytesc                 C   s2   d}|d| j j 7 }| jD ]	}|d|j 7 }q|S )Nz)==== SMB2 Compression Capabilities ====
Flags: %s
zCompression type: %s
)r   r\   r   )r0   r?   r   r   r   r   r^     s
   
z$SMB2CompressionCapabilities.__repr__N)F)	r	   r
   r   r1   r_   rI   r   r:   r^   r   r   r   r   rq     s    

rq   c                   @   r   )SMB2TransportCapabilityFlagsr   r   N)r	   r
   r   r   ACCEPT_TRANSPORT_LEVEL_SECURITYr   r   r   r   r     r   r   c                   @   0   e Zd Zdd Zedd Zdd Zdd Zd	S )
rs   c                 C   s   t j| _d | _d| _d | _d S rv   )r`   rg   ru   rw   r&   r   r/   r   r   r   r1        
z"SMB2TransportCapabilities.__init__c                 C   sr   t  }ttj| dddd|_tj| dddd|_tj| dddd|_ttj| dddd|_	|S r}   )
rq   r`   rE   rF   rG   ru   rw   r&   r   r   rS   r   r   r   r   rI     s   z%SMB2TransportCapabilities.from_bufferc                 C   sh   | j jdddd}t|| _| jjjdddd}|| jjdddd7 }|| jjdddd7 }||7 }|S r}   )r   r:   r;   rw   ru   r=   r&   )r0   r   r?   r   r   r   r:     s   
z"SMB2TransportCapabilities.to_bytesc                 C      d}|d| j  7 }|S )Nz'==== SMB2 Transport Capabilities ====
r   )r   r0   r?   r   r   r   r^        z"SMB2TransportCapabilities.__repr__Nr	   r
   r   r1   r_   rI   r:   r^   r   r   r   r   rs     s    
rs   c                   @   r   )
rr   c                 C   s   t j| _d | _d | _d | _d S rn   )r`   rh   ru   rw   r&   NetNamer/   r   r   r   r1     r   z&SMB2NetnameNegotiateContextID.__init__c                 C   sj   t  }ttj| dddd|_tj| dddd|_tj| dddd|_| |jd|_	|S )Nr   r2   Fr3   r   z	utf-16-le)
rr   r`   rE   rF   rG   ru   rw   r&   decoder   r   r   r   r   rI     s   z)SMB2NetnameNegotiateContextID.from_bufferc                 C   s   | j dS )Nz	utf-18-le)r   encoder/   r   r   r   r:     s   z&SMB2NetnameNegotiateContextID.to_bytesc                 C   r   )Nz-==== SMB2 Netname Negotiate Context ID ====
zNetName: %s
)r   r   r   r   r   r^      r   z&SMB2NetnameNegotiateContextID.__repr__Nr   r   r   r   r   rr     s    
	rr   c                   @   r   )NEGOTIATE_REPLYc                 C   sp   d| _ d | _d | _d | _d | _d | _d| _d| _d| _d | _	d | _
d| _d| _d| _d| _d | _g | _d | _d S )NA   i   r   r6   )r#   r%   DialectRevisionr+   
ServerGuidr'   MaxTransactSizeMaxReadSizeMaxWriteSize
SystemTimeServerStartTimeSecurityBufferOffsetSecurityBufferLengthr*   Bufferr   r-   SecurityBufferr/   r   r   r   r1     s$   
zNEGOTIATE_REPLY.__init__c                 C   s   t t| S rn   )r   rI   rL   BytesIO)bbuffr   r   r   rF     s   zNEGOTIATE_REPLY.from_bytesc                 C   s  t  }tj| dddd|_|jdksJ ttj| dddd|_ttj| dddd|_tj| dddd|_	t
| d|_ttj| dddd|_tj| dddd|_tj| dddd|_tj| dddd|_t| d|_t| d|_tj| dddd|_tj| dddd|_tj| dddd|_|  }|jd	kr| |jtj | |j|_|  }|jtjkrg |_| |jtj t|j	D ]}|jt !|  d|   d }| |tj" q|S )
Nr   r2   Fr3   r   r   r   r   r   )#r   rE   rF   rG   r#   r   r%   r   r   r+   rH   r   r   r'   r   r   r   timestamp2datetimer   r   r   r   r*   rJ   rK   rL   rP   r   r<   r-   rN   rO   rQ   rI   rM   )rS   msgrU   rW   rV   padsizer   r   r   rI   !  s:   
zNEGOTIATE_REPLY.from_bufferc                 C   s  d| _ t| j| _| jjdddd}|| jjdddd7 }|| jjjdddd7 }|| j	jdddd7 }|| j
 7 }|| jjdddd7 }|| jjdddd7 }|| jjdddd7 }|| jjdddd7 }|t| j7 }|t| j7 }|| j jdddd7 }|| jjdddd7 }|| jjdddd7 }|| j7 }| j	dkrt| jD ]*\}}|| 7 }|| j	d krtt|d	 d
\}}|d|d d
 t|  7 }q|S )N   r   r2   Fr3   r   r   r   r   r   r8   )r   r;   r   r   r#   r:   r%   r   r=   r+   r   r'   r   r   r   datetime2timestampr   r   r*   r>   r-   rR   )r0   r?   rC   ngctxrY   rZ   r   r   r   r:   E  s2   

zNEGOTIATE_REPLY.to_bytesc                 C   s   d}|d| j  7 }|dt| j 7 }|d| jj 7 }|d| j 7 }|dt| j 7 }|d| j 7 }|d| j 7 }|d	| j	 7 }|d
| j
  7 }|d| j  7 }|d| j 7 }|d| j 7 }|d| j 7 }|d| j 7 }|S )Nz ==== SMB2 NEGOTIATE REPLY ====
r[   zSecurityMode: %s
zDialectRevision: %s
zServerGuid: %s
zCapabilities: %s
zMaxTransactSize: %s
zMaxReadSize: %s
zMaxWriteSize: %s
zSystemTime: %s
zServerStartTime: %s
zSecurityBufferOffset: %s
zSecurityBufferLength: %s
zBuffer: %s
zNegotiateContextList: %s
)r#   r]   r%   r   r\   r   r'   r   r   r   r   	isoformatr   r   r   r   r-   r   r   r   r   r^   d  s    zNEGOTIATE_REPLY.__repr__N)	r	   r
   r   r1   r_   rF   rI   r:   r^   r   r   r   r   r     s    

#r   )contextvarsr   enumrL   datetimeuuidaiosmb.commons.utils.ts2dtaiosmb.protocol.commonr   .aiosmb.wintypes.dtyp.constrcuted_security.guidIntFlagr   r   r!   Enumr`   rl   rQ   ro   r   rp   r   rt   r   r   rq   IntEnumr   rs   rr   r   r   r   r   r   <module>   s4    v
@66@!