增强的python控制windows命令行程序
之前写过一篇关于python控制命令行的程序:
python控制windows命令行程序
使用之后发现, 对于普通内置dos命令没有问题, 但是对于有些控制台程序没有作用, 比如python程序, 就捕获不到输出信息.
经过查阅相关资料, 发现有些控制台程序需要真正的终端才能够正常运行. windows有一个伪终端的功能可以解决这个问题.
有重新实现了一个版本, 可以捕获任何windows控制台程序了. 实现要点:
- 必须调用windows api函数, 可以通过ctypes来实现
- 捕获的原始内容包括ansi控制字符, 需要特殊处理, 才能够获得正常的控制台文本输出.
附主要源码:(api封装略. ansi_converter如有需要可私信)
import ansi_converter
import Common
import conpyt.kernel32 as api
import ctypes
from ctypes import byref,cast,sizeof
import re
import time
class ConPty():
'''命令行控制台对象'''
_non_value = object()
'''表示没有提供值的标记'''
def __init__(self,cmd:str='cmd.exe',width:int=80,height:int=24) -> None:
'''
cmd: 命令行程序.
width: 控制台宽度(字符数)
height: 控制台高度(字符数)
'''
self._cmd:str = cmd
'''可执行命令行程序'''
self._cmdIn:api.HANDLE = api.HANDLE()
'''命令行输入句柄. HANDLE类型.'''
self._cmdOut:api.HANDLE = api.HANDLE()
'''命令行输出句柄. HANDLE类型.'''
self._height:int = height
'''伪终端字符高度'''
self._hPC:api.HANDLE = api.HANDLE()
'''伪终端PTY句柄'''
self._is_closed:bool = False
'''表示是否资源已经释放'''
self._last_output:str = ''
'''上次输出的文本'''
self._output:str = ''
'''本次输出的文本信息'''
self._prompt:str = r'[a-zA-Z]:\\[^>]+>$'
'''命令行提示符'''
self._ptyIn:api.HANDLE = api.HANDLE()
'''PTY输入句柄. HANDLE类型.'''
self._ptyOut:api.HANDLE = api.HANDLE()
'''PTY输出句柄. HANDLE类型.'''
self._width:int = width
'''伪终端字符宽度'''
def _create_process(self):
'''创建进程'''
self._startupInfoEx = api.STARTUPINFOEX()
self._startupInfoEx.StartupInfo.cb = sizeof(api.STARTUPINFOEX)
lpSize = api.SIZE_T()
api.InitializeProcThreadAttributeList(None, 1, 0, byref(lpSize)) # 第一次, 获取结构的尺寸
mem = api.HeapAlloc(api.GetProcessHeap(), 0, lpSize.value)
self._mem = mem # 记录分配的内存, 后续需要释放
self._startupInfoEx.lpAttributeList = cast(mem,api.LPSTRUCT)
api.InitializeProcThreadAttributeList(self._startupInfoEx.lpAttributeList, 1, 0, byref(lpSize))
PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE = 0x00020016
api.UpdateProcThreadAttribute(self._startupInfoEx.lpAttributeList,
0,
cast(PROC_THREAD_ATTRIBUTE_PSEUDOCONSOLE,api.DWORD_PTR),
self._hPC,
sizeof(self._hPC),
None,
None
) # 更新进程属性信息
self._lpProcessInformation = api.PROCESS_INFORMATION()
EXTENDED_STARTUPINFO_PRESENT = 0x00080000
hr = api.CreateProcess(None,
self._cmd,
None,
None,
False,
EXTENDED_STARTUPINFO_PRESENT,
None,
None,
byref(self._startupInfoEx.StartupInfo),
byref(self._lpProcessInformation))
if hr==0: raise Exception('20240326_2215: 创建进程失败.')
api.WaitForSingleObject(self._lpProcessInformation.hThread, 500)
def _create_pty(self):
'''创建pty伪终端'''
api.CreatePipe(byref(self._ptyIn), byref(self._cmdIn), None, 0) # 创建输入管道
api.CreatePipe(byref(self._cmdOut), byref(self._ptyOut), None, 0) # 创建输出管道
hr = api.CreatePseudoConsole(api.COORD(self._width,self._height), self._ptyIn, self._ptyOut, 0, byref(self._hPC))
if hr!=0: raise Exception(f'20240319_204401: CreatePseudoConsole failed with error code {hr}')
def close(self):
'''释放相关资源. 采用倒序释放(先申请, 后释放)的方式.'''
api.CloseHandle(self._lpProcessInformation.hThread)
api.CloseHandle(self._lpProcessInformation.hProcess)
api.DeleteProcThreadAttributeList(self._startupInfoEx.lpAttributeList)
api.HeapFree(api.GetProcessHeap(), 0, self._mem)
api.ClosePseudoConsole(self._hPC)
api.CloseHandle(self._ptyOut)
api.CloseHandle(self._cmdOut)
api.CloseHandle(self._cmdIn)
api.CloseHandle(self._ptyIn)
self._is_closed = True
def __del__(self):
'''对象释放时调用. 会清除相关资源'''
if not self._is_closed: self.close()
def _expect(self,expect:str):
'''等待期望的输出'''
if expect==None: return # 如果没有期望输出, 则直接返回
if expect==self._non_value: expect = self._prompt # 如果用户没有输入参数值, 则使用缺省参数
while True:
if re.search(expect,self._output,re.MULTILINE): break
time.sleep(0.1)
self._output = '' # 一旦匹配成功, 则清空输出, 以免影响后续的匹配
def get_last_output(self)->str:
'''获取上一次的输出'''
return self._last_output
@Common.run_in_thread
def _read(self):
'''循环读取输出'''
def _has_new_data()->bool:
'''判断是否有新数据'''
nonlocal lpBuffer
lpNumberOfBytesRead = api.DWORD()
lpNumberOfBytesAvail = api.DWORD()
lpBytesLeftInMessage = api.DWORD()
hr = api.PeekNamedPipe(self._cmdOut,
lpBuffer,
MAX_READ,
byref(lpNumberOfBytesRead),
byref(lpNumberOfBytesAvail),
byref(lpBytesLeftInMessage))
return not hr == 0x0 and lpNumberOfBytesAvail.value > 0
ansi_converter.set_con_size(self._width, self._height) # 先设置虚拟屏幕的尺寸
MAX_READ = 1024*8
lpBuffer = ctypes.create_string_buffer(MAX_READ)
lpNumberOfBytesRead = api.DWORD()
TIME_OUT = 0.2 # 超时时间, 秒
TIME_INTERVAL = 0.01 # 轮询间隔, 秒
while True:
# 先读取数据, 如果没有数据, 则阻塞在这里. 这样就避免了轮询操作
ret = api.ReadFile(self._cmdOut,lpBuffer,MAX_READ,byref(lpNumberOfBytesRead),None)
if ret==0: break # 读取失败, 则中断
_output_data:bytes = lpBuffer.raw[:lpNumberOfBytesRead.value] # 保存读取的数据
# 下面进行一小波读取. 避免数据不正常截断...
begin_time = time.time() # 这一波数据的起始时间
while True: # 循环检查是否还有新数据, 直到超时为止
if _has_new_data(): # 如果有新数据, 则读取它, 追加到之前的结果中
api.ReadFile(self._cmdOut,lpBuffer,MAX_READ,byref(lpNumberOfBytesRead),None)
_output_data += lpBuffer.raw[:lpNumberOfBytesRead.value]
begin_time = time.time() # 重置起始时间
continue
# 如果没有新数据, 检查是否超时. 如果没有超时, 则继续轮询
if time.time()-begin_time<TIME_OUT:
time.sleep(TIME_INTERVAL)
continue
# 超时了, 把当前这一波数据输出到屏幕上...
decoded_string = _output_data.decode('utf-8')
plain = ansi_converter.render_ansi_text(decoded_string)
print(plain,end='')
self._output = self._last_output = plain # 设置本次输出的文本信息
break # 完成输出, 等待下一波数据的到来
def start(self,expect:str|None=_non_value):
'''
启动命令.
expect: 正则表达式表示的期望输出. 如果输入None表示不判断期望输出. 缺省使用prompt作为期望输出.
'''
self._create_pty() # 创建PTY伪终端
self._create_process() # 必须先创建进程, 否则读取线程, 读取会失败
self._read() # 启动读取线程
self._expect(expect)
def writeline(self,line,expect:str|None=_non_value):
'''
向pty写入一行数据. 会自动添加回车换行符
expect: 正则表达式表示的期望输出. 如果输入None表示不判断期望输出. 缺省使用prompt作为期望输出.
'''
line += '\r\n'
lpBuffer = ctypes.create_string_buffer(line.encode('utf-8'))
lpNumberOfBytesWritten = api.DWORD()
api.WriteFile(self._cmdIn,lpBuffer,sizeof(lpBuffer),byref(lpNumberOfBytesWritten),None)
self._expect(expect)
if __name__ == '__main__':
pty = ConPty()
pty.start()
pty.writeline('python',expect=r'>>>$')
pty.writeline('print("胡中强")',expect=r'>>>$')
pty.close()
