a
    %gh                     @   s  d dl Z d dlZd dlZd dlmZmZmZmZmZm	Z	m
Z
mZmZ d dlmZ edjZdZdZdZdZejjZdd	d
Zdd ZG dd dejZegej_eej_eegej_eej_eeej_ee egej!_eej!_eeej!_ee egej"_eej"_eeej"_g ej#_eej#_eeegej$_eej$_eeej$_eee	egej%_eej%_eeej%_e j&dd Z'dd Z(dd Z)dddZ*dS )    N)	BOOLCHARDWORDHANDLELONGLPWSTRMAX_PATHPDWORDULONG)SHELL_NAMES   z      i   c                    s    fdd}|S )Nc                    s   |  krt  | S N)ctypesWinError)retfuncargs	error_val ^/mounts/lovelace/software/anaconda3/envs/metaDMG/lib/python3.9/site-packages/shellingham/nt.pycheck   s    z_check_handle.<locals>.checkr   )r   r   r   r   r   _check_handle   s    r   c                    s    fdd}|S )Nc                    s*   | rdS t  }| krdS t |d S )NTF)r   ZGetLastErrorr   )r   r   r   codeexpectedr   r   r   (   s    z_check_expected.<locals>.checkr   )r   r   r   r   r   _check_expected'   s    r   c                   @   sV   e Zd Zdefdefdefdeefdefdefdefdefd	efd
ee	 ff
Z
dS )ProcessEntry32dwSizeZcntUsageth32ProcessIDZth32DefaultHeapIDZth32ModuleIDZ
cntThreadsth32ParentProcessIDZpcPriClassBaseZdwFlags	szExeFileN)__name__
__module____qualname__r   r   POINTERr
   r   r   r   Z_fields_r   r   r   r   r    3   s   
r    c              	   o   s4   | |i |}z|V  W t | nt | 0 d S r   )kernel32CloseHandle)fr   kwargshandler   r   r   _handleg   s    r.   c                  c   sl   t j} t| tdF}t }t||_t ||}|rJ|V  t 	||}q2W d    n1 s^0    Y  d S )Nr   )
r)   CreateToolhelp32Snapshotr.   TH32CS_SNAPPROCESSr    r   Zsizeofr!   Process32FirstProcess32Next)r+   Zsnapentryr   r   r   r   _iter_processesp   s    r4   c                 C   s@   t t}td|j}t| d||r,|jS | jd9  _qd S )N r   r   )r   r   r   Zcreate_unicode_buffervaluer)   QueryFullProcessImageNameW)prochsizeZ	path_buffr   r   r   _get_full_path{   s
    r:   
   c           	   	   C   s   dd t  D }| pt } td|d D ]}z||  \}}W n tyV   Y  qY n0 t|trn|dd}|dd 	 }|t
vr|} q*t}ttj|d|  }|t|fW  d      S 1 s0    Y  q*d S )Nc                 S   s   i | ]}|j |j|jfqS r   )r"   r#   r$   ).0procr   r   r   
<dictcomp>   s   zget_shell.<locals>.<dictcomp>r      mbcsreplace.)r4   osgetpidrangeKeyError
isinstancebytesdecode
rpartitionlowerr   !PROCESS_QUERY_LIMITED_INFORMATIONr.   r)   OpenProcessr:   )	pid	max_depthZproc_map_Zppid
executablenamekeyr8   r   r   r   	get_shell   s&    
0rT   )r   )Nr;   )+
contextlibr   rC   Zctypes.wintypesr   r   r   r   r   r   r   r	   r
   Zshellingham._corer   r6   ZINVALID_HANDLE_VALUEZERROR_NO_MORE_FILESZERROR_INSUFFICIENT_BUFFERr0   rL   Zwindllr)   r   r   Z	Structurer    r*   argtypesrestyper/   Zerrcheckr(   r1   r2   ZGetCurrentProcessIdrM   r7   contextmanagerr.   r4   r:   rT   r   r   r   r   <module>   s^   ,

	

	