mirror of https://github.com/dnomd343/ProxyC
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
203 lines
9.3 KiB
203 lines
9.3 KiB
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
import os
|
|
import copy
|
|
import time
|
|
import ctypes
|
|
import signal
|
|
from Basis.Logger import logging
|
|
from Basis.Functions import genFlag
|
|
from subprocess import Popen, STDOUT, DEVNULL
|
|
|
|
libcPaths = [
|
|
'/usr/lib/libc.so.6', # CentOS
|
|
'/usr/lib64/libc.so.6',
|
|
'/lib/libc.musl-i386.so.1', # Alpine
|
|
'/lib/libc.musl-x86_64.so.1',
|
|
'/lib/libc.musl-aarch64.so.1',
|
|
'/lib/i386-linux-gnu/libc.so.6', # Debian / Ubuntu
|
|
'/lib/x86_64-linux-gnu/libc.so.6',
|
|
'/lib/aarch64-linux-gnu/libc.so.6',
|
|
]
|
|
|
|
libcPath = None
|
|
for libc in libcPaths:
|
|
if os.path.exists(libc): # try to locate libc.so
|
|
libcPath = libc
|
|
break
|
|
if libcPath is None: # lost libc.so -> unable to utilize prctl
|
|
logging.warning('libc.so not found')
|
|
else:
|
|
logging.info('libc.so -> ' + str(libcPath))
|
|
|
|
|
|
class Process(object):
|
|
""" Manage a sub process and it's file.
|
|
|
|
Arguments:
|
|
cmd: Command list, which use to start the program.
|
|
|
|
env: Environment variables for sub process.
|
|
|
|
file: dict or list or None, include path and content option.
|
|
|
|
workDir: A directory for storing log files and configuration files.
|
|
|
|
taskId: Task ID, defaults to 12 random characters length.
|
|
|
|
isStart: Start the process after class init complete.
|
|
|
|
Attributes:
|
|
id, cmd, env, file, workDir, output
|
|
"""
|
|
output = None # sub process output if capture is True
|
|
__capture = None # capture the sub process output or not
|
|
__logfile = None # the log file which sub process output into STDOUT and STDERR
|
|
__process = None # CompletedProcess object of subprocess module
|
|
|
|
@staticmethod
|
|
def __preExec() -> None:
|
|
ctypes.CDLL(libcPath).prctl(1, signal.SIGTERM) # sub process killed when father process exit
|
|
os.setpgrp() # new process group
|
|
|
|
def __checkWorkDir(self) -> None: # check if the working directory is normal
|
|
if os.path.isdir(self.workDir):
|
|
return
|
|
logging.warning('[%s] Work directory %s not exist' % (self.id, self.workDir))
|
|
try:
|
|
os.makedirs(self.workDir) # just like `mkdir -p ...`
|
|
logging.info('[%s] New directory -> %s' % (self.id, self.workDir))
|
|
except:
|
|
if os.path.exists(self.workDir):
|
|
logging.error('[%s] %s already exist but not folder' % (self.id, self.workDir))
|
|
else:
|
|
logging.error('[%s] Unable to create new folder -> %s' % (self.id, self.workDir))
|
|
raise RuntimeError('working directory error') # fatal error
|
|
|
|
def __killProcess(self, killSignal: int) -> None:
|
|
try:
|
|
pgid = os.getpgid(self.__process.pid) # progress group id
|
|
os.killpg(pgid, killSignal) # kill sub process group
|
|
logging.debug('[%s] Send kill signal to PGID %i' % (self.id, pgid))
|
|
except:
|
|
logging.warning('[%s] Failed to get PGID of sub process (PID = %i)' % (self.id, self.__process.pid))
|
|
|
|
def __deleteFile(self, filePath: str) -> None:
|
|
if not os.path.isfile(filePath): # file not found (or not a file)
|
|
logging.warning('[%s] File %s not found' % (self.id, filePath))
|
|
return
|
|
try:
|
|
os.remove(filePath) # remove config file
|
|
logging.debug('[%s] File %s deleted successfully' % (self.id, filePath))
|
|
except:
|
|
logging.error('[%s] Unable to delete file %s' % (self.id, filePath))
|
|
|
|
def __init__(self, workDir: str, taskId: str = '', isStart: bool = True,
|
|
cmd: str or list or None = None, env: dict or None = None, file: dict or list or None = None) -> None:
|
|
self.id = genFlag(length = 12) if taskId == '' else taskId
|
|
self.workDir = workDir
|
|
self.env = copy.copy(env) # depth = 1
|
|
self.cmd = copy.copy([cmd] if type(cmd) == str else cmd) # depth = 1
|
|
self.file = copy.deepcopy([file] if type(file) == dict else file) # depth = 2
|
|
self.__checkWorkDir() # ensure the working direction is normal
|
|
logging.debug('[%s] Process command -> %s (%s)' % (self.id, self.cmd, self.env))
|
|
if self.file is not None:
|
|
if len(self.file) > 1:
|
|
logging.debug('[%s] Manage %i files' % (self.id, len(self.file)))
|
|
for file in self.file: # traverse all files
|
|
if not isStart: # don't print log twice
|
|
logging.debug('[%s] File %s -> %s' % (self.id, file['path'], file['content']))
|
|
if isStart:
|
|
self.start()
|
|
|
|
def setCmd(self, cmd: str or list) -> None:
|
|
self.cmd = copy.copy([cmd] if type(cmd) == str else cmd)
|
|
logging.info('[%s] Process setting command -> %s' % (self.id, self.cmd))
|
|
|
|
def setEnv(self, env: dict or None) -> None:
|
|
self.env = copy.copy(env)
|
|
logging.info('[%s] Process setting environ -> %s' % (self.id, self.env))
|
|
|
|
def setFile(self, file: dict or list or None) -> None:
|
|
self.file = copy.deepcopy([file] if type(file) == dict else file)
|
|
if self.file is None:
|
|
logging.info('[%s] Process setting file -> None' % self.id)
|
|
return
|
|
for file in self.file: # traverse all files
|
|
logging.info('[%s] Process setting file %s -> %s' % (self.id, file['path'], file['content']))
|
|
|
|
def start(self, isCapture: bool = True) -> None:
|
|
self.__capture = isCapture
|
|
logging.debug('[%s] Process ready to start (%s)' % (self.id, (
|
|
'with output capture' if self.__capture else 'without output capture'
|
|
)))
|
|
if self.cmd is None: # ERROR CASE
|
|
logging.error('[%s] Process miss start command' % self.id)
|
|
raise RuntimeError('miss start command')
|
|
if self.__process is not None and self.__process.poll() is None: # ERROR CASE
|
|
logging.error('[%s] Sub process is still running' % self.id)
|
|
raise RuntimeError('sub process is still running')
|
|
if self.env is not None and 'PATH' not in self.env and '/' not in self.cmd[0]: # WARNING CASE
|
|
logging.warning('[%s] Executable file in relative path but miss PATH in environ' % self.id)
|
|
if self.file is not None: # create and write file contents
|
|
for file in self.file:
|
|
with open(file['path'], 'w', encoding = 'utf-8') as fileObject: # save file content
|
|
fileObject.write(file['content'])
|
|
logging.debug('[%s] File %s -> %s' % (self.id, file['path'], file['content']))
|
|
if self.__capture: # with output capture
|
|
self.__logfile = os.path.join(self.workDir, self.id + '.log')
|
|
logging.debug('[%s] Output capture file -> %s' % (self.id, self.__logfile))
|
|
stdout = open(self.__logfile, 'w', encoding = 'utf-8')
|
|
stderr = STDOUT # combine the stderr with stdout
|
|
else: # discard all the output of sub process
|
|
stdout = DEVNULL
|
|
stderr = DEVNULL
|
|
self.__process = Popen(
|
|
self.cmd, env = self.env,
|
|
stdout = stdout, stderr = stderr,
|
|
preexec_fn = None if libcPath is None else Process.__preExec
|
|
)
|
|
logging.info('[%s] Process running -> PID = %i' % (self.id, self.__process.pid))
|
|
|
|
def signal(self, signalNum: int) -> None: # send specified signal to sub process
|
|
try:
|
|
signalName = signal.Signals(signalNum).name
|
|
except:
|
|
signalName = 'unknown'
|
|
logging.info('[%s] Send signal -> %i (%s)' % (self.id, signalNum, signalName))
|
|
self.__process.send_signal(signalNum)
|
|
|
|
def status(self) -> bool: # check if the sub process is still running
|
|
status = self.__process.poll() is None
|
|
logging.debug('[%s] Check status -> %s' % (self.id, 'running' if status else 'exit'))
|
|
return status
|
|
|
|
def wait(self, timeout: int or None = None) -> None: # blocking wait sub process
|
|
logging.info('[%s] Process wait -> timeout = %s' % (self.id, str(timeout)))
|
|
try:
|
|
self.__process.wait(timeout = timeout)
|
|
logging.info('[%s] Process wait timeout -> sub process exit' % self.id)
|
|
except:
|
|
logging.info('[%s] Process wait timeout -> sub process still running' % self.id)
|
|
|
|
def quit(self, isForce: bool = False, waitTime: int = 50) -> None: # wait 50ms in default
|
|
killSignal = signal.SIGKILL if isForce else signal.SIGTERM # 9 -> force kill / 15 -> terminate signal
|
|
logging.debug('[%s] Kill signal = %i (%s)' % (self.id, killSignal, signal.Signals(killSignal).name))
|
|
self.__killProcess(killSignal)
|
|
time.sleep(waitTime / 1000) # sleep (ms -> s)
|
|
while self.__process.poll() is None: # confirm sub process exit
|
|
self.__killProcess(killSignal)
|
|
time.sleep(waitTime / 1000)
|
|
logging.info('[%s] Process terminated -> PID = %i' % (self.id, self.__process.pid))
|
|
if self.__capture:
|
|
try:
|
|
with open(self.__logfile, 'r', encoding = 'utf-8') as fileObject: # read sub process output
|
|
self.output = fileObject.read()
|
|
logging.debug('[%s] Process output -> length = %s' % (self.id, len(self.output)))
|
|
self.__deleteFile(self.__logfile)
|
|
except:
|
|
logging.error('[%s] Failed to read capture log file -> %s' % (self.id, self.__logfile))
|
|
if self.file is not None: # with config file
|
|
for file in self.file:
|
|
self.__deleteFile(file['path'])
|
|
|