mirror of https://github.com/openspug/spug
74 lines
2.2 KiB
Python
74 lines
2.2 KiB
Python
from libs.utils import str_decode
|
|
import subprocess
|
|
import re
|
|
import os
|
|
|
|
|
|
class Executor:
|
|
def __init__(self, default_env=None):
|
|
self.regex = re.compile(r'(?<!echo )Spug EOF 2108111926 (-?\d+)[\r\n]?')
|
|
self.default_env = default_env
|
|
self.task = None
|
|
|
|
@property
|
|
def pid(self):
|
|
if self.task:
|
|
return self.task.pid
|
|
return None
|
|
|
|
def get_envs(self):
|
|
envs = {}
|
|
_, output = self.exec_command('env | grep SPUG_')
|
|
if output:
|
|
for item in output.splitlines():
|
|
if '=' in item:
|
|
key, val = item.split('=', 1)
|
|
envs[key] = val
|
|
return envs
|
|
|
|
def exec_command_with_stream(self, command):
|
|
command += '\necho Spug EOF 2108111926 $?\n'
|
|
self.task.stdin.write(command.encode())
|
|
self.task.stdin.flush()
|
|
exit_code, message, = -1, b''
|
|
while True:
|
|
output = self.task.stdout.read(1)
|
|
if not output:
|
|
exit_code, message = self.task.wait(), ''
|
|
break
|
|
if output in (b'\r', b'\n'):
|
|
message += b'\r\n' if output == b'\n' else b'\r'
|
|
message = str_decode(message)
|
|
match = self.regex.search(message)
|
|
if match:
|
|
exit_code = int(match.group(1))
|
|
message = message[:match.start()]
|
|
break
|
|
yield exit_code, message
|
|
message = b''
|
|
else:
|
|
message += output
|
|
yield exit_code, message
|
|
|
|
def exec_command(self, command):
|
|
exit_code, message = -1, ''
|
|
for exit_code, line in self.exec_command_with_stream(command):
|
|
message += line
|
|
return exit_code, message
|
|
|
|
def __enter__(self):
|
|
self.task = subprocess.Popen(
|
|
'bash',
|
|
env=self.default_env,
|
|
shell=True,
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
preexec_fn=os.setsid
|
|
)
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
if self.task:
|
|
self.task.kill()
|