goat_monitor/goat_monitor.py

94 lines
2.6 KiB
Python
Raw Normal View History

# %% imports
import subprocess
from pathlib import Path
from typing import List
import click
import gotify
2024-09-14 12:45:24 -06:00
import numpy as np
import toml
__version__ = "v0.0.0"
# %% commands
@click.command()
@click.argument("command", nargs=-1, required=False, type=str)
@click.option(
"--config",
default="~/.config/goat_monitor.toml",
type=click.Path(path_type=Path),
help="Use a .toml configuration file",
)
2024-09-14 12:45:24 -06:00
@click.option(
"--retries",
default=0,
type=int,
help="Number of times to try re-running failed command. -1 to retry indefinitely until success",
)
def wrap(command: List[str], config: Path, retries: int):
"""Wrap an arbitrary command with gotify notifications"""
2024-09-13 21:00:44 -06:00
# read settings first
with open(config, "r") as f:
settings = toml.load(f)
2024-09-14 12:45:24 -06:00
# gotify configuration
2024-09-13 21:00:44 -06:00
url = settings["server"]
2024-09-14 12:45:24 -06:00
app_token = settings["app_token"]
with gotify.Gotify(base_url=url, app_token=app_token) as gotify_connection:
# this is just to test configuration.
# I don't know how long a gotify connection will stay up so rather than thinking about it
# I'm just creating a new one whenever I need to send a message
pass
2024-09-13 21:00:44 -06:00
2024-09-14 12:45:24 -06:00
if retries == -1:
retries = np.inf
if retries < 0:
raise ValueError("Invalid number of retries specified")
for attempt in range(retries + 1):
# run the command
result = subprocess.run(
2024-09-14 12:45:24 -06:00
" ".join(command),
shell=True,
stderr=subprocess.STDOUT,
stdout=subprocess.PIPE,
encoding="utf-8",
)
2024-09-13 21:00:44 -06:00
# TODO: print lines real time as the subprocess runs
print(result.stdout, end="")
2024-09-14 12:45:24 -06:00
if result.returncode:
# failed
title = f"Command failed with exit code {result.returncode}"
if attempt < retries:
title += f" - Retrying (attempt {attempt+1}/{retries})"
else:
title += " - Aborting"
else:
title = "Command succeeded"
MAX_LINES = 20
lines = result.stdout.splitlines()
if len(lines) > MAX_LINES:
lines = lines[-MAX_LINES:]
message = (
"Command:\n"
+ " ".join(command)
+ f'\n\nResult{ " (truncated)" if len(lines) > MAX_LINES else ""}:\n'
+ "\n".join(lines)
)
with gotify.Gotify(base_url=url, app_token=app_token) as gotify_connection:
gotify_connection.create_message(message=message, title=title)
2024-09-13 21:00:44 -06:00
return result.returncode
# %% main
if __name__ == "__main__":
wrap()