108 lines
3.2 KiB
Python
108 lines
3.2 KiB
Python
|
import logging
|
||
|
|
import shutil
|
||
|
|
import sys
|
||
|
|
import os
|
||
|
|
import signal
|
||
|
|
import subprocess
|
||
|
|
import argparse
|
||
|
|
import rich
|
||
|
|
from rich.logging import RichHandler
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import antidrift.client
|
||
|
|
import antidrift.daemon
|
||
|
|
from antidrift.config import Config
|
||
|
|
|
||
|
|
from dbus.mainloop.glib import DBusGMainLoop
|
||
|
|
|
||
|
|
DBusGMainLoop(set_as_default=True)
|
||
|
|
signal.signal(signal.SIGINT, signal.SIG_DFL)
|
||
|
|
|
||
|
|
|
||
|
|
def get_args():
|
||
|
|
parser = argparse.ArgumentParser(description="AntiDrift CLI.")
|
||
|
|
parser.add_argument("--daemon", action="store_true", help="run daemon")
|
||
|
|
parser.add_argument("--status", action="store_true", help="get status from daemon")
|
||
|
|
parser.add_argument("--tailf", action="store_true", help="tail -f log file")
|
||
|
|
parser.add_argument(
|
||
|
|
"--start", metavar="whiteblock", nargs="+", help="start whiteblocks"
|
||
|
|
)
|
||
|
|
parser.add_argument("--stop", action="store_true", help="stop session")
|
||
|
|
parser.add_argument("--schedule", metavar="blackblock", help="schedule blackblock")
|
||
|
|
args = parser.parse_args()
|
||
|
|
return args
|
||
|
|
|
||
|
|
|
||
|
|
def init_logging(log_file: Path, dev_mode: bool = False):
|
||
|
|
class DuplicateFilter(logging.Filter):
|
||
|
|
def filter(self, record) -> bool:
|
||
|
|
current_log = (record.module, record.levelno, record.msg)
|
||
|
|
if current_log != getattr(self, "last_log", None):
|
||
|
|
self.last_log = current_log
|
||
|
|
return True
|
||
|
|
return False
|
||
|
|
|
||
|
|
if dev_mode:
|
||
|
|
format_str = "%(message)s" # RichHandler will handle the formatting
|
||
|
|
logging.basicConfig(
|
||
|
|
level=logging.DEBUG,
|
||
|
|
format=format_str,
|
||
|
|
datefmt="%a %H:%M:%S",
|
||
|
|
handlers=[RichHandler(rich_tracebacks=True, markup=True)],
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
format_str = (
|
||
|
|
"[bold pale_green3]%(asctime)s[/bold pale_green3] | "
|
||
|
|
"[light_steel_blue]%(levelname)-8s[/light_steel_blue] | "
|
||
|
|
"%(message)s"
|
||
|
|
)
|
||
|
|
logging.basicConfig(
|
||
|
|
filename=log_file,
|
||
|
|
format=format_str,
|
||
|
|
datefmt="%a %H:%M:%S",
|
||
|
|
encoding="utf-8",
|
||
|
|
level=logging.DEBUG,
|
||
|
|
)
|
||
|
|
|
||
|
|
logger = logging.getLogger()
|
||
|
|
logger.addFilter(DuplicateFilter())
|
||
|
|
|
||
|
|
|
||
|
|
def check_for_xdotool():
|
||
|
|
"""Check if xdotool is in path and exit if not"""
|
||
|
|
result = shutil.which("xdotool")
|
||
|
|
if not result:
|
||
|
|
logging.critical("Please install xdotool")
|
||
|
|
sys.exit(1)
|
||
|
|
|
||
|
|
|
||
|
|
def main_daemon(config):
|
||
|
|
if os.geteuid() == 0:
|
||
|
|
newpid = os.fork()
|
||
|
|
if newpid == 0:
|
||
|
|
init_logging(config.daemon_log_file)
|
||
|
|
antidrift.daemon.run(config)
|
||
|
|
else:
|
||
|
|
if sys.argv[0] == "antidrift":
|
||
|
|
cmd = ["sudo", "antidrift", "--daemon"]
|
||
|
|
subprocess.Popen(cmd)
|
||
|
|
else:
|
||
|
|
init_logging(config.daemon_log_file, dev_mode=True)
|
||
|
|
logging.warning("[red]Running in development mode.[/red]")
|
||
|
|
antidrift.daemon.run(config)
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
"""Main routine that dispatches to client or daemon"""
|
||
|
|
config = Config.load(os.path.expanduser("~/.config/antidrift.yaml"))
|
||
|
|
check_for_xdotool()
|
||
|
|
args = get_args()
|
||
|
|
if args.daemon:
|
||
|
|
main_daemon(config)
|
||
|
|
else:
|
||
|
|
antidrift.client.run(args, config)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|