o
    j:                     @   s   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 d dl	m
Z
 G dd dZdd	 Zdd
dZdd Zdd ZedkrKd dlmZ e  dS dS )    N)Fields)NegotiateFlags)Version)AVPairsc                   @   sR   e Zd Zdd Zedd Zedd Zeddd	Zd
d Zdd Z	dd Z
dS )NTLMChallengec                 C   sF   d| _ d| _d | _d | _d | _d| _d | _d | _d | _d | _	d | _
d S )Ns   NTLMSSP    s           )	SignatureMessageTypeTargetNameFieldsr   ServerChallengeReservedTargetInfoFieldsr   Payload
TargetName
TargetInfoself r   /root/aizidognhua/tmp/workspace/projects/ec89d86c-575f-41c9-af57-ac45cbdbf775/venv/lib/python3.10/site-packages/asyauth/protocols/ntlm/messages/challenge.py__init__   s   
zNTLMChallenge.__init__c                 C   s   t t| S N)r   from_bufferioBytesIO)bbuffr   r   r   
from_bytes   s   zNTLMChallenge.from_bytesc                 C   s8  t  }| d|_tj| dddd|_t| |_t	tj| dddd|_	| d|_
| d|_t| |_|j	t	j@ rGt| |_|  }|  |_|jjdkr~| |jjtj | |jj}z|d|_W n ty}   |d|_Y nw |jjdkr| |jjtj | |jj}t||_|S )	N      littleF	byteordersignedr   utf-16lezutf-8)r   readr   intr   r	   r   r   r
   r   r   r   r   NEGOTIATE_VERSIONr   tellr   lengthseekoffsetr   SEEK_SETdecoder   UnicodeDecodeErrorr   r   )bufftcurrPosraw_datar   r   r   r      s2   
zNTLMChallenge.from_bufferNc           	      C   sr  | du r	t d} t| dkrtdt }| |_||_|du r%td}n	t|ts.t|}|dur7|tj	O }||_d}|durD|d7 }|dur_|
d}||_tt|||_|t|7 }n	d|_td||_|dur|| }||_tt|||_nt |_|j }tt|||_d|_|dur| j|
d7  _|dur| j| 7  _|S | j|j 7  _|S )	a  
		Construct a new NTLMChallenge message.
		
		Args:
			challenge: 8-byte server challenge (defaults to random)
			targetName: Target name string (optional)
			targetInfo: AVPairs object for target info (optional) 
			version: Version object (optional)
			flags: NegotiateFlags object (required if version is provided)
		
		Returns:
			NTLMChallenge: Constructed challenge message
		Nr   z(Server challenge must be exactly 8 bytesr   0   r"        )osurandomlen
ValueErrorr   r   r   r   
isinstancer%   encoder   r   r
   to_bytesr   r   r   r   )		challenge
targetName
targetInfoversionflagsr.   payload_postarget_name_datatarget_info_datar   r   r   	constructB   sN   





zNTLMChallenge.constructc                 C   s   | j d}| j }| j}|| jjdddd7 }|| j 7 }|| jjdddd7 }|| j7 }|| j	7 }|| j
 7 }| jrF|| j 7 }|| j7 }|S )Nr"   r   r   Fr   )r   r9   r   r:   r   r	   r
   r   r   r   r   r   r   )r   tntir-   r   r   r   r:      s   



zNTLMChallenge.to_bytesc                 C   sb   d}|dt | j 7 }|dt | j 7 }|dt | j 7 }|dt | j 7 }|dt | j 7 }|S )Nz== NTLMChallenge ==
zSignature      : %s
zMessageType    : %s
zServerChallenge: %s
zTargetName     : %s
zTargetInfo     : %s
)reprr   r	   r   r   r   )r   r.   r   r   r   __repr__   s   zNTLMChallenge.__repr__c                 C   s   t |  dS )Nascii)base64	b64encoder:   r+   r   r   r   r   toBase64   s   zNTLMChallenge.toBase64)NNNNN)__name__
__module____qualname__r   staticmethodr   r   rC   r:   rG   rK   r   r   r   r   r      s    

"L	r   c                   C   s   t   t  t  d S r   )test_reconstruttest_constructtest_templater   r   r   r   test   s   
rS   c                 C   s   zddl m} W n ty   dd }Y nw td | s"td}n| }t|}tt| |	 }td t|| td t|| ||ksMJ d S )	Nr   hexdumpc                 S   s   |   S r   )hex)datar   r   r   rU      s   z test_reconstrut.<locals>.hexdumpz=== reconstruct===d  4e544c4d53535000020000000800080038000000158289e2a7314a557bdb11bf000000000000000072007200400000000a0063450000000f540045005300540002000800540045005300540001001200570049004e003200300031003900410044000400120074006500730074002e0063006f007200700003002600570049004e003200300031003900410044002e0074006500730074002e0063006f007200700007000800aec600bfc5fdd40100000000z====== reconstructed ====z====== original ====)
asyauth.utils.hexdumprU   ImportErrorprintbytesfromhexr   r   rF   r:   )rW   rU   challenge_test_datar;   challenge_test_data_verifyr   r   r   rP      s"   
rP   c                   C   s   d S r   r   r   r   r   r   rR      s   rR   c                  C   s   t d t } t d d}tj|d}|j|ksJ t d tjdd}|jdks,J t d d	d
lm}m} | }d||j< tj|d}|j	|j dksPJ t d tj|d|d}|
 }	t|	}
|
j|kslJ |
jdkssJ t d t d d S )Nz=== test_construct ===zBasic construct test passeds   )r;   zCustom challenge test passedDOMAIN)r<   zTarget name test passedr   )r   
AVPAIRType)r=   zTarget info test passedTEST)r;   r<   r=   zRoundtrip test passedzAll construct tests passed!)r[   r   rC   r   r   (asyauth.protocols.ntlm.structures.avpairr   ra   MsvAvNbDomainNamer   r:   r   )
challenge1custom_challenge
challenge2
challenge3r   ra   target_info
challenge4
challenge5rW   
challenge6r   r   r   rQ      s6   

rQ   __main__rT   r   )r4   r   rI   (asyauth.protocols.ntlm.structures.fieldsr   1asyauth.protocols.ntlm.structures.negotiate_flagsr   )asyauth.protocols.ntlm.structures.versionr   rc   r   r   rS   rP   rR   rQ   rL   rY   rU   r   r   r   r   <module>   s"     #
(
