import abc import base64 import json import locale import os import subprocess import sys import time from subprocess import CREATE_NO_WINDOW _blockInput = None _messageBox = None if sys.platform == 'win32': import ctypes from ctypes import wintypes import win32ui # import win32con _messageBox = win32ui.MessageBox _blockInput = ctypes.windll.user32.BlockInput _blockInput.argtypes = [wintypes.BOOL] _blockInput.restype = wintypes.BOOL def block_input(): if _blockInput: _blockInput(True) def unblock_input(): if _blockInput: _blockInput(False) def decode_content(content: bytes) -> str: for encoding_name in ['utf-8', 'gbk', 'gb2312']: try: return content.decode(encoding_name) except Exception as e: print(e) encoding_name = locale.getpreferredencoding() return content.decode(encoding_name) def notify_err_message(msg): if _messageBox: _messageBox(msg, 'Error') def check_pid_alive(pid) -> bool: # tasklist /fi "PID eq 508" /fo csv # '"映像名称","PID","会话名 ","会话# ","内存使用 "\r\n"wininit.exe","508","Services","0","6,920 K"\r\n' try: csv_ret = subprocess.check_output(["tasklist", "/fi", f'PID eq {pid}', "/fo", "csv"], creationflags=CREATE_NO_WINDOW) content = decode_content(csv_ret) content_list = content.strip().split("\r\n") if len(content_list) != 2: print("check pid {} ret invalid: {}".format(pid, content)) return False ret_pid = content_list[1].split(",")[1].strip('"') return str(pid) == ret_pid except Exception as e: print("check pid {} err: {}".format(pid, e)) return False def wait_pid(pid): while 1: time.sleep(5) ok = check_pid_alive(pid) if not ok: print("pid {} is not alive".format(pid)) break class DictObj: def __init__(self, in_dict: dict): assert isinstance(in_dict, dict) for key, val in in_dict.items(): if isinstance(val, (list, tuple)): setattr(self, key, [DictObj(x) if isinstance(x, dict) else x for x in val]) else: setattr(self, key, DictObj(val) if isinstance(val, dict) else val) class User(DictObj): id: str name: str username: str class Specific(DictObj): # web autofill: str username_selector: str password_selector: str submit_selector: str script: list # database db_name: str use_ssl: str class Secret(DictObj): client_key: str class Category(DictObj): value: str label: str class Protocol(DictObj): id: str name: str port: int class Asset(DictObj): id: str name: str address: str protocols: list[Protocol] category: Category spec_info: Specific secret_info: Secret def get_protocol_port(self, protocol): for item in self.protocols: if item.name == protocol: return item.port return None class LabelValue(DictObj): label: str value: str class Account(DictObj): id: str name: str username: str secret: str privileged: bool secret_type: LabelValue class Platform(DictObj): charset: str name: str charset: LabelValue type: LabelValue class Manifest(DictObj): name: str version: str path: str exec_type: str connect_type: str protocols: list[str] def get_manifest_data() -> dict: current_dir = os.path.dirname(__file__) manifest_file = os.path.join(current_dir, 'manifest.json') try: with open(manifest_file, "r", encoding='utf8') as f: return json.load(f) except Exception as e: print(e) return {} def read_app_manifest(app_dir) -> dict: main_json_file = os.path.join(app_dir, "manifest.json") if not os.path.exists(main_json_file): return {} with open(main_json_file, 'r', encoding='utf8') as f: return json.load(f) def convert_base64_to_dict(base64_str: str) -> dict: try: data_json = base64.decodebytes(base64_str.encode('utf-8')).decode('utf-8') return json.loads(data_json) except Exception as e: print(e) return {} class BaseApplication(abc.ABC): def __init__(self, *args, **kwargs): self.app_name = kwargs.get('app_name', '') self.protocol = kwargs.get('protocol', '') self.manifest = Manifest(kwargs.get('manifest', {})) self.user = User(kwargs.get('user', {})) self.asset = Asset(kwargs.get('asset', {})) self.account = Account(kwargs.get('account', {})) self.platform = Platform(kwargs.get('platform', {})) @abc.abstractmethod def run(self): raise NotImplementedError('run') @abc.abstractmethod def wait(self): raise NotImplementedError('wait')