o
    Žõ±j2>  ã                   @   s¾   d dl Z d dlZd dlZd dlZd dlmZ d dlT d dlT d dlT d dl	m
Z
 d dlmZ d dlmZ d dlmZ g d¢ZG d	d
„ d
ƒZG dd„ dƒZdd„ Zdd„ Zedkr]eƒ  dS dS )é    N)Úlogger)Ú*)ÚSMBConnectionFactory)Ú
SMBMachine)ÚUniversalEncoder)Útqdm)	ÚtargetÚ	target_idÚotypeÚpathÚcreationtimeÚsizeÚsizefmtÚsddlÚerrc                   @   s8   e Zd Zdd„ Zdd„ Zdd„ Zdd„ Zefd	d
„ZdS )ÚEnumResultFinalc                 C   sâ   || _ || _|| _|| _|| _d | _d | _d | _d | _d | _	d | _
| jdv rm| j j
| _
| jdks:| jdks:| jdkr[| jdksD| jdkrI| j j| _| j j| _| jd u rUdn| j ¡ | _	| jdkro| j j| _t| jƒ| _d S d S d S )N)ÚdirÚfileÚsharer   r   r   Ú )Úobjr
   r   r   r	   Úcreation_timer   Úsize_fmtÚsecurity_descriptorÚsecurity_descriptor_sddlÚunc_pathÚto_sddlÚ
sizeof_fmt)Úselfr   r
   r   r   r	   © r   ú�/root/aizidognhua/tmp/workspace/projects/ec89d86c-575f-41c9-af57-ac45cbdbf775/venv/lib/python3.10/site-packages/aiosmb/examples/smbshareenum.pyÚ__init__   s.   





özEnumResultFinal.__init__c                 C   sª   | j d urd| j| j f S | jdkr d| j| j| j| j| jf S | jdkr/d| j| j| jf S | jdkr<d| j| jf S | jdkrSd	| jj| jj	t
| jjƒ| jjf S d
S )Nz[E] %s | %sr   z[F] %s | %s | %s | %s | %sr   z[D] %s | %s | %sr   z[S] %s | %sÚprogressz[P][%s/%s][%s] %sz[UNK])r   r   r
   r   r   r   r   r   Útotal_targetsÚtotal_finishedÚstrÚgens_finishedÚcurrent_finished©r   r   r   r    Ú__str__2   s   




$zEnumResultFinal.__str__c                 C   s.   | j | j| j| j| j| j| j| j| j| j	dœ
S )N)
r   r	   r   r   r   r   Úsecuritydescriptorr   r
   r   )
r   r	   r   r   r   r   r   r   r
   r   r(   r   r   r    Úto_dictE   s   özEnumResultFinal.to_dictc                 C   s   |   ¡ }|d= tj|td�S )Nr*   )Úcls)r+   ÚjsonÚdumpsr   )r   Úddr   r   r    Úto_jsonS   s   zEnumResultFinal.to_jsonc                    s2   | j dkrdS |  ¡ ‰ ‡ fdd„|D ƒ}d |¡S )Nr"   r   c                    s   g | ]}t ˆ | ƒ‘qS r   ©r%   ©Ú.0Úx©r/   r   r    Ú
<listcomp>\   s    z*EnumResultFinal.to_tsv.<locals>.<listcomp>ú	)r
   r+   Újoin)r   ÚhdrsÚdatar   r5   r    Úto_tsvX   s
   

zEnumResultFinal.to_tsvN)	Ú__name__Ú
__module__Ú__qualname__r!   r)   r+   r0   ÚENUMRESFINAL_TSV_HDRr;   r   r   r   r    r      s    r   c                   @   sr   e Zd Zdddddddddddddg g g ddfdd	„Zd
d„ Zdd„ Zdd„ Zdd„ Zdd„ Zdd„ Z	dd„ Z
dS )ÚSMBFileEnumé
   é   FNTr%   é   c                 C   sÔ   g | _ || _t|tƒrt |¡| _|| _|| _|| _|| _	g | _
d | _|| _|| _|| _|| _|| _|	| _|
| _|| _|| _|| _|| _|| _|| _|| _d| _d| _d| _d| _d| _d| _d| _ d| _!i | _"d S )NFr   )#Útarget_gensÚsmb_mgrÚ
isinstancer%   r   Úfrom_urlÚworker_countÚtask_qÚres_qÚdepthÚworkersÚresult_processing_taskÚenum_urlÚout_fileÚ	show_pbarÚ	max_itemsÚmax_runtimeÚfetch_share_sdÚfetch_dir_sdÚfetch_file_sdÚoutput_typeÚexclude_shareÚexclude_dirÚexclude_targetÚext_result_qÚconnection_timeoutÚ_SMBFileEnum__gens_finishedÚ_SMBFileEnum__total_targetsÚ_SMBFileEnum__total_finishedÚ_SMBFileEnum__total_sizeÚ_SMBFileEnum__total_sharesÚ_SMBFileEnum__total_dirsÚ_SMBFileEnum__total_filesÚ_SMBFileEnum__total_errorsÚ_SMBFileEnum__current_targets)r   Úsmb_urlrH   rK   rN   rO   rP   rQ   rR   rS   rT   rU   rI   rJ   rV   rW   rX   rY   rZ   r[   r   r   r    r!   a   sB   

zSMBFileEnum.__init__c                 Ã   sÊ  �zÒzm| j  |¡}|4 I d H šT tj| ¡ | jd�I d H \}}|d ur%|‚t|ƒ}|j| j| j	| j
| j| j| j| jd�2 z3 d H W \}}}t|||||fƒ}	| j |	¡I d H  q<6 W d   ƒI d H  n1 I d H siw   Y  W n> tjy‹   Y W | j t||d tjd�¡I d H  d S  ty­ }
 z| j t||d |
tjd�¡I d H  W Y d }
~
nd }
~
ww W | j t||d tjd�¡I d H  d S W | j t||d tjd�¡I d H  d S | j t||d tjd�¡I d H  w )N©Útimeout)rK   Ú
maxentriesrS   rT   rU   rW   rX   ©Ústatus©Úerrorrj   )rE   Úcreate_connection_newtargetÚasyncioÚwait_forÚloginr[   r   Úenum_all_recursivelyrK   rQ   rS   rT   rU   rW   rX   Ú
EnumResultrJ   ÚputÚCancelledErrorÚEnumResultStatusÚFINISHEDÚ	ExceptionÚERROR)r   Útidr   Ú
connectionÚ_r   Úmachiner   r
   ÚerÚer   r   r    Ú
__executor†   s2   €8þ(ú€
$ý.€ÿ€$þFzSMBFileEnum.__executorc                 Ã   s8  �z|	 | j  ¡ I d H }|d u rW d S |\}}ztj|  ||¡| jd�I d H  W nU tjy3   Y W d S  tjyf } z'| j 	t
||d |tjd�¡I d H  | j 	t
||d tjd�¡I d H  W Y d }~qd }~w ty| } zt d¡ W Y d }~qd }~ww q tjyˆ   Y d S  ty› } z|W  Y d }~S d }~ww )NTrf   rk   ri   Úworker)rI   Úgetrn   ro   Ú_SMBFileEnum__executorrR   rt   ÚTimeoutErrorrJ   rs   rr   ru   rx   rv   rw   r   Ú	exception)r   Úindatary   r   r~   r   r   r    r€   š   s8   €"" €
€þò€ÿzSMBFileEnum.workerc                 Ã   s4  ��zìd }| j du rUi }tdddd�|d< tdddd�|d	< td
ddd�|d< tdddd�|d< tddddd�|d< tdddd�|d< tdddd�|d< tdddd�|d< g }d}	 �ze| jdu r~| j du r~|d jd u r~| j|d _|D ]}||  ¡  qu| jd ur…g }t|ƒdks”|�r2| jd u �r2d}| jd kr¦d! 	d"d#„ |D ƒ¡}n+| jd$kr¹|D ]
}|d%| 
¡  7 }q­n| jd&krÌ|D ]
}|d%| ¡  7 }qÀnd! 	|¡}| jd u�rt| jd'dd(��/}z| |¡ W n   | d!¡D ]}z	| |d! ¡ W qð   t d)¡ Y qðY W d   ƒ n	1 �sw   Y  nt|ƒ | j du �r.|D ]	}||  ¡  �q$g }d}|�r@t |  ¡ ¡ W W d S ztj| j ¡ dd*�I d H }	W n- tj�y|   | j du �rj|D ]	}||  ¡  �q`| j| jk�ry| jdu �ryd}Y W qYw |	jtjk�rØ|  jd7  _| j du �r˜|d  d¡ t| j| j| j|	j ƒ}
| jd u�rº| j !t"|
d+d |	j |	j#ƒ¡I d H  | $t"|
d+d |	j |	j#ƒ¡ | j| jk�rØ| jdu �rØd}W qY|	j%d u�r¤|	j%\}
}}| jd u�rû| j !t"|
|||	j |	j#ƒ¡I d H  | $t"|
|||	j |	j#ƒ¡ |d u�r‹|d,k�r9|  j&d7  _&t'|
j(t)ƒdu �r8|  j*|
j(7  _*| j du �r8|d  |
j(¡ n|d-k�rF|  j+d7  _+n|d.k�rR|  j,d7  _,| j du �r‹|d-k�re|d  d¡ n&|d,k�rr|d  d¡ n|d.k�r|d	  d¡ n|dk�r‹|d  d¡ |d u�r¤|  j-d7  _-| j du �r¤|d  d¡ |	jtj.k�r¿|  j-d7  _-| j du �r¿|d  d¡ W n- tj/�yÍ   Y W d S  t0�yí } zt 1d/¡ t |  ¡ ¡ W Y d }~W d S d }~ww qZ tj/�yú   Y d S  t0�y } zt 1d0¡ t |  ¡ ¡ W Y d }~d S d }~ww )1NTzTargets     r   r   )ÚdescÚunitÚpositionÚtargetszShares      é   ÚshareszDirs        é   ÚdirszFiles       rB   ÚfileszFiles (size)ÚBé   )r†   r‡   Ú
unit_scalerˆ   ÚfilesizezMaxed       rC   ÚmaxedzEnum Errors é   Ú
enumerrorszConn Errors é   Ú
connerrorsFiè  r%   z
c                 S   s   g | ]}t |ƒ‘qS r   r1   r2   r   r   r    r6   Ð   s    z1SMBFileEnum.result_processing.<locals>.<listcomp>Útsvz%s
r-   za+)ÚnewlinezEncoding issue, skipping a linerf   r"   r   r   r   zresult_processing innerÚresult_processing)2rP   r   r\   Útotalr]   ÚrefreshrZ   ÚlenrV   r8   r;   r0   rO   ÚopenÚwriteÚsplitr   ÚdebugÚprintrn   Úcreate_taskÚ	terminatero   rJ   r�   rƒ   r^   rj   ru   rv   ÚupdateÚEnumProgressr   rs   r   r	   ÚappendÚresultrb   rF   r   Úintr_   ra   r`   rc   rx   rt   rw   r„   )r   ÚpbarÚ
out_bufferÚ
final_iterÚkeyÚout_dataÚresÚfÚliner}   r   r
   r   r~   r   r   r    rš   ³   s  €
"


ÿ
ÿ

û€ü€ù	""

€






€
€ý�t
€þzSMBFileEnum.result_processingc                 Ã   s2   �| j D ]}| ¡  q| jd ur| j ¡  d S d S ©N)rL   ÚcancelrM   )r   r€   r   r   r    r¤   =  s   €


ÿzSMBFileEnum.terminatec              
   Ã   sœ   �z6| j d u rt | j¡| _ t |  ¡ ¡| _| jd u r"t | j¡| _t| jƒD ]}| j	 
t |  ¡ ¡¡ q'W dS  tyM } z
d |fW  Y d }~S d }~ww )N©TN)rJ   rn   ÚQueuerH   r£   rš   rM   rI   ÚrangerL   r§   r€   rw   )r   r{   r~   r   r   r    ÚsetupC  s   €

€ÿzSMBFileEnum.setupc                 Ã   sð   �| j du r!|  jd7  _| j tt ¡ ƒ| j ¡  	¡ f¡I d H  | j
D ]=}| ¡ 2 z53 d H W \}}}|d ur?td| ƒ  n"|| jv rEq*|  jd7  _| j ||f¡I d H  t d¡I d H  q*6 q$t| jƒD ]}| j d ¡I d H  qgd| _d S )NTrŠ   zTarget gen error! %sr   )rN   r]   rI   rs   r%   ÚuuidÚuuid4rE   Ú
get_targetÚget_hostname_or_iprD   Úgenerater¢   rY   rn   Úsleepr¶   rH   r\   )r   Ú
target_genÚuidr   r   r{   r   r   r    Ú__generate_targetsR  s&   €
(

ö€
zSMBFileEnum.__generate_targetsc                 Ã   s  �zqz?|   ¡ I d H \}}|d ur|‚t |  ¡ ¡}tj| jŽ I d H  | jI d H  W W | jd ur@| j t	d dd d d ƒ¡I d H  dS dS  t
yr } z%t d¡ d |fW  Y d }~W | jd urm| j t	d dd d d ƒ¡I d H  S S d }~ww | jd urˆ| j t	d dd d d ƒ¡I d H  w w )NÚfinishedr´   Úrun)r·   rn   r£   Ú_SMBFileEnum__generate_targetsÚgatherrL   rM   rZ   rs   r   rw   r   r„   )r   r{   r   Úgen_taskr~   r   r   r    rÂ   h  s.   €
"ÿü

 ÿ€ü
 ÿzSMBFileEnum.run)r<   r=   r>   r!   r‚   r€   rš   r¤   r·   rÃ   rÂ   r   r   r   r    r@   `   s    .% r@   c                  Ã   sú  �dd l } dd l}d}| jd| j|d�}|jddddd� |jd	td
dd� |jddtddd� |jdddd� |jddddd� |jdtd dd� |jdtd dd� |jddd� |jd dd!d� |jd"dd#d� |jd$dd%d� |jd&dd'd� |jd(dd)d� |jd*dd+d� |jd,d-d.d/d� |jd0d1d.d2d� |jd3td4d5d� |jd6d7d8d9� | ¡ }|jd:kr¼t 	t
j¡ |jd;krØtd<ƒ t 	d:¡ t ¡  d=¡ t
jt
jd>� d?}|jd=u rád@}|jd=u rèdA}|j}g }|jd urõ|j}g }|jd urÿ|j}t||j|j|j|j|j|j|j|j||j|||j dB�}	g }
t!|j"ƒdk�r5|j#d=u �r5|	j$ %t&|j#ƒ¡ n%|j"D ]!}zt'|dCƒ}| (¡  |	j$ %t)|ƒ¡ W �q8   |
 %|¡ Y �q8t!|
ƒdk�ri|	j$ %t&|
ƒ¡ t!|	j$ƒdk�rtd=|	_*|	 +¡ I d H  d S )DNr   z£
Output legend:
    [S] Share
    [D] Dictionary
    [F] File
    [E] Error
    [M] Maxed (max items limit reached for directory)
    [P] Progress (current/total)
zSMB Share enumerator)ÚdescriptionÚformatter_classÚepilogz-vz	--verboseÚcount)ÚactionÚdefaultz--depthrB   z"Recursion depth, -1 means infinite)ÚtyperË   Úhelpz-wz--smb-worker-countéd   zParallell countz-oz
--out-filezOutput file path.)rÍ   z-sz--stdinÚ
store_truezRead targets from stdin)rÊ   rÍ   z--max-itemsz>Stop enumeration of a directory after N items were discovered.z--max-runtimez*Stop enumeration of a host after N secondsz--urlz—Connection URL base, target can be set to anything. Owerrides all parameter based connection settings! Example: "smb2+ntlm-password://TEST\victim@test"z
--progresszShow progress barz	--sharesdzFetch share security descriptorz--dirsdz#Fetch directory security descriptorz--filesdzFetch file security descriptorz--jsonzOutput in JSON formatz--tsvz,Output in TSV format. (TAB Separated Values)z--esz--exclude-sharer§   z"Exclude shares with name specifiedz--edz--exclude-dirz'Exclude directories with name specifiedz--connection-timeoutrC   zConnection timeoutr‰   r   z5Hostname or IP address or file with a list of targets)ÚnargsrÍ   rŠ   rŒ   zsetting deepdebugT)Úlevelr%   r-   r˜   )rH   rK   rO   rP   rQ   rS   rT   rU   rV   rR   rW   rX   r[   Úr),ÚargparseÚsysÚArgumentParserÚRawDescriptionHelpFormatterÚadd_argumentr©   Ú
parse_argsÚverboser   ÚsetLevelÚloggingÚDEBUGr¢   rn   Úget_event_loopÚ	set_debugÚbasicConfigr-   r˜   ÚurlÚesÚedr@   Úsmb_worker_countrK   rO   r"   rQ   ÚsharesdÚdirsdÚfilesdrR   r[   r�   r‰   ÚstdinrD   r§   ÚListTargetGenrž   ÚcloseÚFileTargetGenrN   rÂ   )rÓ   rÔ   rÈ   ÚparserÚargsrV   re   rW   rX   Ú
enumeratorÚnotfiler   r°   r   r   r    Úamainz  s”   €







ò

rï   c                   C   s   t  tƒ ¡ d S r²   )rn   rÂ   rï   r   r   r   r    ÚmainÝ  s   rð   Ú__main__)rn   r¸   rÛ   r-   Úaiosmbr   Ú&aiosmb.examples.scancommons.targetgensÚ$aiosmb.examples.scancommons.internalÚ!aiosmb.examples.scancommons.utilsÚ!aiosmb.commons.connection.factoryr   Ú!aiosmb.commons.interfaces.machiner   Ú!aiosmb.commons.utils.univeraljsonr   r   r?   r   r@   rï   rð   r<   r   r   r   r    Ú<module>   s,   J  c
ÿ