148 lines
4.7 KiB
Python
Executable File
148 lines
4.7 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import os.path
|
|
import rx
|
|
import rx.operators
|
|
import rx.scheduler
|
|
import sys
|
|
import threading
|
|
import typing
|
|
|
|
from decimal import Decimal
|
|
from solana.publickey import PublicKey
|
|
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
import mango # nopep8
|
|
|
|
parser = argparse.ArgumentParser(
|
|
description="Watches one or many accounts (via a websocket) and sends a notification if the SOL balance falls below the --minimum-sol-balance threshold."
|
|
)
|
|
mango.ContextBuilder.add_command_line_parameters(parser)
|
|
mango.Wallet.add_command_line_parameters(parser)
|
|
parser.add_argument(
|
|
"--named-address",
|
|
type=str,
|
|
required=True,
|
|
action="append",
|
|
default=[],
|
|
help="Name and address of the Solana account to watch, separated by a colon",
|
|
)
|
|
parser.add_argument(
|
|
"--minimum-sol-balance",
|
|
type=Decimal,
|
|
default=Decimal("0.1"),
|
|
help="the minimum SOL balance required for the alert. A SOL balance less than this value will trigger a nofitication.",
|
|
)
|
|
parser.add_argument(
|
|
"--timer-limit",
|
|
type=int,
|
|
default=(60 * 60),
|
|
help="notifications for an account will be sent at most once per timer-limit seconds, and accounts will be polled once per timer-limit seconds irrespective of websocket activity",
|
|
)
|
|
parser.add_argument(
|
|
"--notify",
|
|
type=mango.parse_notification_target,
|
|
action="append",
|
|
default=[],
|
|
help="The notification target for low balance events",
|
|
)
|
|
parser.add_argument(
|
|
"--notify-events",
|
|
type=mango.parse_notification_target,
|
|
action="append",
|
|
default=[],
|
|
help="The notification target for startup events",
|
|
)
|
|
args: argparse.Namespace = mango.parse_args(parser)
|
|
|
|
notify_balance: mango.NotificationTarget = mango.CompoundNotificationTarget(args.notify)
|
|
notify_event: mango.NotificationTarget = mango.CompoundNotificationTarget(
|
|
args.notify_events
|
|
)
|
|
|
|
|
|
def notifier(name: str) -> typing.Callable[[mango.AccountInfo], None]:
|
|
def notify(account_info: mango.AccountInfo) -> None:
|
|
report = f'Account "{name} [{account_info.address}]" on {context.client.cluster_name} has only {account_info.sols} SOL, which is below the minimum required balance of {args.minimum_sol_balance} SOL.'
|
|
notify_balance.send(f"[{args.name}] {report}")
|
|
mango.output(f"Notification sent: {report}")
|
|
|
|
return notify
|
|
|
|
|
|
def account_fails_balance_check(account_info: mango.AccountInfo) -> bool:
|
|
return bool(account_info.sols < args.minimum_sol_balance)
|
|
|
|
|
|
def log_account(account_info: mango.AccountInfo) -> mango.AccountInfo:
|
|
logging.info(f"Checking account {account_info.address} - {account_info.sols} SOL.")
|
|
return account_info
|
|
|
|
|
|
def add_subscription_for_parameter(
|
|
context: mango.Context,
|
|
manager: mango.WebSocketSubscriptionManager,
|
|
health_check: mango.HealthCheck,
|
|
timer_limit: int,
|
|
name_and_address: str,
|
|
) -> None:
|
|
name, address_str = name_and_address.split(":")
|
|
address = PublicKey(address_str)
|
|
|
|
immediate = mango.AccountInfo.load(context, address)
|
|
if immediate is None:
|
|
raise Exception(f"No account '{name}' at {address_str}.")
|
|
|
|
account_subscription = mango.WebSocketAccountSubscription(
|
|
context, address, lambda account_info: account_info
|
|
)
|
|
manager.add(account_subscription)
|
|
|
|
on_change = account_subscription.publisher.pipe(rx.operators.start_with(immediate))
|
|
on_timer = rx.interval(timer_limit).pipe(
|
|
rx.operators.map(lambda _: mango.AccountInfo.load(context, address))
|
|
)
|
|
rx.merge(on_change, on_timer).pipe(
|
|
rx.operators.observe_on(context.create_thread_pool_scheduler()),
|
|
rx.operators.map(log_account),
|
|
rx.operators.filter(account_fails_balance_check),
|
|
rx.operators.throttle_first(timer_limit),
|
|
rx.operators.catch(mango.observable_pipeline_error_reporter),
|
|
rx.operators.retry(),
|
|
).subscribe(notifier(name))
|
|
|
|
|
|
context = mango.ContextBuilder.from_command_line_parameters(args)
|
|
|
|
disposer = mango.DisposePropagator()
|
|
manager = mango.SharedWebSocketSubscriptionManager(context)
|
|
disposer.add_disposable(manager)
|
|
health_check = mango.HealthCheck()
|
|
disposer.add_disposable(health_check)
|
|
|
|
# Need a nice way of ensuring pongs from all subscriptions
|
|
health_check.add("ws_pong", manager.pong)
|
|
|
|
for name_and_address in args.named_address:
|
|
add_subscription_for_parameter(
|
|
context, manager, health_check, args.timer_limit, name_and_address
|
|
)
|
|
|
|
manager.open()
|
|
|
|
notify_event.send(f"[{args.name}] Starting using context:\n{context}")
|
|
|
|
# Wait - don't exit. Exiting will be handled by signals/interrupts.
|
|
waiter = threading.Event()
|
|
try:
|
|
waiter.wait()
|
|
except:
|
|
pass
|
|
|
|
logging.info("Shutting down...")
|
|
disposer.dispose()
|
|
logging.info("Shutdown complete.")
|