SteamAppMonitor.py

# -*- coding: utf-8 -*-
import ctypes
import logging
import logging.handlers
import os
import re
import sys
import threading
import time

from PIL import Image, ImageDraw
from pystray import Icon, Menu, MenuItem
import click
import darkdetect as dd
import psutil

# --- 設定項目 ---
APP_ID = None                   # SteamのAppID
MEMORY_THRESHOLD_MB = 1024      # メモリしきい値 (MB)
CHECK_INTERVAL = 60             # 監視間隔 (秒)
PROCESS_NAME = None             # SteamApp の実行ファイル名
# ----------------
PreferredAppMode = {
    'Light': 0,
    'Dark': 1,
}
# https://github.com/moses-palmer/pystray/issues/130
ctypes.windll['uxtheme.dll'][135](PreferredAppMode[dd.theme()])


def find_exec_file(APP_ID):
    STEAM_LOG_FILE = r'C:\Program Files (x86)\Steam\logs\gameprocess_log.txt'
    FIND_PATTERN = re.compile(rf'AppID {APP_ID} adding PID \d+ as a tracked process "(?P<path>"?.*?"?)"')
    exec_file = str()

    with open(STEAM_LOG_FILE, encoding='utf-8') as fd:
        m = re.search(FIND_PATTERN, fd.read())
        if m:
            exec_file = os.path.basename(m.groups('path')[0])

    return exec_file.strip('"')


class SteamMonitorApp:
    def __init__(self):
        self.stop_event = threading.Event()
        self.monitor_thread = None
        self.app = None
        self.logger = None

    def get_total_memory_usage(self, process_name):
        """指定した名前のプロセス（および子プロセス）の合計メモリ使用量(MB)を取得"""
        total_mem = 0
        for proc in psutil.process_iter(['name', 'memory_info']):
            try:
                if proc.info['name'].lower() == process_name.lower():
                    total_mem += proc.info['memory_info'].rss
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                continue
        return total_mem / (1024 * 1024)

    def restart_steam_app(self):
        """アプリを終了させてSteam経由で再起動"""
        # プロセスを強制終了
        for proc in psutil.process_iter(['name']):
            if proc.info['name'].lower() == PROCESS_NAME.lower():
                proc.kill()

        # SteamのURLスキームを使用して起動
        os.startfile(f"steam://run/{APP_ID}")
        self.logger.info(f'Threshold exceeded. Restarting AppID: {APP_ID}')

    def monitor_loop(self):
        while not self.stop_event.is_set():
            start_time = time.time()
            mem_usage = self.get_total_memory_usage(PROCESS_NAME)
            title = f'{APP_ID}: {PROCESS_NAME} {int(mem_usage)} / {MEMORY_THRESHOLD_MB}'
            self.logger.info(title)
            self.app.title = title
            self.app.update_menu()
            elapsed = time.time() - start_time

            if mem_usage > MEMORY_THRESHOLD_MB:
                self.restart_steam_app()
                # 再起動直後に連続判定されないよう少し待機
                time.sleep(30)

            time.sleep(CHECK_INTERVAL - elapsed)

    def create_image(self):
        width, height = 64, 64
        image = Image.new('RGB', (width, height), (255, 255, 255))
        dc = ImageDraw.Draw(image)
        dc.ellipse([10, 10, 54, 54], fill=(0, 120, 215))
        return image

    def stopApp(self, icon, item):
        self.stop_event.set()
        self.app.stop()

    def runApp(self):
        self.monitor_thread = threading.Thread(target=self.monitor_loop, daemon=True)
        self.monitor_thread.start()

        # システムトレイアイコンの設定
        self.app = Icon("SteamMonitor")
        self.app.menu = Menu(
            MenuItem(f'AppID: {APP_ID} を監視中...', lambda: None, enabled=False),
            MenuItem('終了', self.stopApp)
        )
        self.app.icon = self.create_image()
        self.app.run()


@click.command(help='SteamAppMonitor')
@click.option('-t', '--threshold', type=int, default=1024)
@click.option('-i', '--interval', type=int, default=60)
@click.argument('appid', type=int, nargs=1, required=True)
def run(appid, threshold, interval):
    global APP_ID, MEMORY_THRESHOLD_MB, CHECK_INTERVAL, PROCESS_NAME

    APP_ID = appid
    MEMORY_THRESHOLD_MB = threshold
    CHECK_INTERVAL = interval

    PROCESS_NAME = find_exec_file(APP_ID)
    if not PROCESS_NAME:
        print(f'No such APP_ID {APP_ID} entried', file=sys.stderr)
        exit(1)
    else:
        # logger settings
        logging.basicConfig(
            format='%(asctime)s [%(levelname)s] %(message)s',
            handlers=[
                logging.handlers.RotatingFileHandler(f'log_{PROCESS_NAME.replace(".exe", "")}.log', encoding='utf-8', maxBytes=1000000, backupCount=0),
                logging.StreamHandler(),
            ],
            datefmt='%Y/%m/%d %X'
        )
        app = SteamMonitorApp()
        app.logger = logging.getLogger(PROCESS_NAME)
        app.logger.setLevel(logging.DEBUG)
        app.runApp()


if __name__ == '__main__':
    run()