def setup_client():
logging.info("Setting up client")
client = ActivityWatchClient("aw-fake-client", testing=True)
client.client_hostname = hostname
eventtype = "currentwindow"
client.create_bucket(window_bucket_name, eventtype)
eventtype = "afkstatus"
client.create_bucket(afk_bucket_name, eventtype)
client.connect()
return client
def delete_prev_buckets():
logging.info("Deleting old buckets")
client = ActivityWatchClient("aw-fake-client", testing=True)
client.client_hostname = hostname
client.connect()
try:
client.delete_bucket(window_bucket_name)
except HTTPError:
pass
try:
client.delete_bucket(afk_bucket_name)
except HTTPError:
pass
class ClientTest(unittest.TestCase):
def setUp(self):
self.client = ActivityWatchClient("unittest", testing=True)
self.client.setup_bucket("test", "testevents")
self.client.connect()
def test_send_event(self):
self.client.send_event("test",
Event(timestamp=datetime.now(), label="test"))
def test_send_events(self):
self.client.send_events("test", [
Event(timestamp=datetime.now(), label="test"),
Event(timestamp=datetime.now(), label="test2"),
])
def setup_client() -> ActivityWatchClient:
logger.info("Setting up client")
# Default is to run in testing mode, can be run in prod mode if set to exactly 'false'
testing = os.getenv("AW_TESTING", "true").lower() not in ["false"]
if testing:
logger.info(
"Using testing parameters (set the env var AW_TESTING to false to run in prod mode)"
)
client = ActivityWatchClient(client_name, testing=testing)
client.client_hostname = hostname
buckets = client.get_buckets()
logger.info("Deleting old buckets")
buckets_all = [
bucket_afk,
bucket_window,
bucket_browser_chrome,
bucket_browser_firefox,
]
if not testing:
ans = input(
f"Running in prod, are you sure you want to delete all existing buckets?\n{buckets_all}\nAre you sure? (y/N)"
)
if ans != "y":
print("Exiting")
sys.exit(0)
for bucket in [
bucket_window,
bucket_afk,
bucket_browser_chrome,
bucket_browser_firefox,
]:
if bucket in buckets:
client.delete_bucket(bucket, force=True)
client.create_bucket(bucket_window, "currentwindow")
client.create_bucket(bucket_afk, "afkstatus")
client.create_bucket(bucket_browser_chrome, "web.tab.current")
client.create_bucket(bucket_browser_firefox, "web.tab.current")
client.connect()
return client
class AFKWatcher:
def __init__(self, testing=False, settings=None):
# Read settings from config
configsection = "aw-watcher-afk" if not testing else "aw-watcher-afk-testing"
self.settings = Settings(watcher_config[configsection])
self.client = ActivityWatchClient("aw-watcher-afk", testing=testing)
self.bucketname = "{}_{}".format(self.client.client_name, self.client.client_hostname)
eventtype = "afkstatus"
self.client.setup_bucket(self.bucketname, eventtype)
self.client.connect()
def set_state(self, status, duration, timestamp=None):
data = {"status": status}
if timestamp is None:
timestamp = self.now
e = Event(data=data, timestamp=timestamp, duration=duration)
self.client.heartbeat(self.bucketname, e, pulsetime=self.settings.timeout, queued=True)
def ping(self, afk):
data = {"status": "afk" if afk else "not-afk"}
e = Event(data=data, timestamp=self.now)
self.client.heartbeat(self.bucketname, e, pulsetime=self.settings.timeout, queued=True)
def run(self):
# TODO: All usage of last_input can probably be replaced the self.seconds_since_last_input equivalent
logger.info("afkwatcher started")
""" Initialization """
sleep(1)
""" Init variables """
self.afk = False
self.now = datetime.now(timezone.utc)
self.last_check = self.now
self.seconds_since_last_input = 0
""" Start afk checking loop """
while True:
try:
self.last_check = self.now
self.now = datetime.now(timezone.utc)
self.seconds_since_last_input = get_seconds_since_last_input()
self.timedelta_since_last_input = timedelta(seconds=self.seconds_since_last_input)
self.last_input = self.now - self.timedelta_since_last_input
logger.debug("Time since last input: {}".format(self.timedelta_since_last_input))
# If program is not allowed to run for more than polling_interval+10s it will assume that the computer has gone into suspend/hibernation
if self.now > self.last_check + timedelta(seconds=10 + self.settings.polling_interval):
logger.info("Woke up from suspend/hibernation")
time_since_last_check = self.now - self.last_check
self.set_state("hibernating", timedelta(seconds=time_since_last_check.total_seconds()), self.last_check)
# If no longer AFK
elif self.afk and self.seconds_since_last_input < self.settings.timeout:
logger.info("No longer AFK")
self.ping(self.afk) # End afk period
self.afk = False
self.set_state("not-afk", timedelta())
# If becomes AFK
elif not self.afk and self.seconds_since_last_input > self.settings.timeout:
logger.info("Became AFK")
self.afk = True
self.set_state("afk", self.timedelta_since_last_input, self.last_input)
# Send a heartbeat if no state change was made
else:
if self.afk:
self.ping(self.afk)
elif self.seconds_since_last_input < self.settings.polling_interval:
self.ping(self.afk)
sleep(self.settings.polling_interval)
except KeyboardInterrupt:
logger.info("afkwatcher stopped by keyboard interrupt")
self.ping(self.afk)
break
class MessageHandler:
def __init__(self,
testing=False,
send_commands=True,
send_heartbeats=True):
# Create client
self._client = ActivityWatchClient(client_id, testing=testing)
self._client.connect()
self._init_buckets()
# Settings
self.pulsetime = 10
self.send_commands = send_commands
self.send_heartbeats = send_heartbeats
# Initialize the EventQueue
self._event_queue = EventQueue(callback=self._handle_event,
time_buffer=(self.pulsetime / 2))
self._event_handlers = {
'preopen': self._preopen,
'preexec': self._preexec,
'precmd': self._precmd,
'preclose': self._preclose
}
self._terminal_sessions = {}
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self._client.disconnect()
def _init_buckets(self):
"""Set self._buckets and create these buckets if not existing"""
self._buckets = {
'commands': {
'id': "{}-commands_{}".format(client_id,
self._client.hostname),
'event_type': 'app.terminal.command'
},
'activity': {
'id': "{}-activity_{}".format(client_id,
self._client.hostname),
'event_type': 'app.terminal.activity'
}
}
# Create buckets if not existing
for key, bucket in self._buckets.items():
logger.debug("Creating bucket: {}".format(bucket['id']))
self._client.create_bucket(bucket['id'],
bucket['event_type'],
queued=True)
def update_event_queue(self):
self._event_queue.update()
def handle_fifo_message(self, message):
for line in message.split('\n'):
if not len(line):
continue
cli_args = shlex.split(line)
args, unknown_args = parser_base.parse_known_args(cli_args)
# Make sure that events are called in the right order
# (based on args.timestamp) by passing it to the event queue
# The event queue will trigger the callback when the time_buffer
# is exceeded
logger.debug("adding event to event_queue: {}".format(cli_args))
self._event_queue.add_event(cli_args, args.timestamp)
def _handle_event(self, cli_args):
logger.debug("handling event: {}".format(cli_args))
args, unknown_args = parser_base.parse_known_args(cli_args)
# Store terminal session if not already existing
pid = args.pid
if pid not in self._terminal_sessions:
self._terminal_sessions[pid] = TerminalSessionData(pid)
if self.send_commands:
if args.event not in self._event_handlers:
logger.error("Unknown event: {}".format(args.event))
else:
self._event_handlers[args.event](cli_args)
if self.send_heartbeats:
self._heartbeat(cli_args)
def _preopen(self, cli_args: list) -> None:
"""Handle terminal creation"""
pass
def _preexec(self, cli_args: list) -> None:
"""Send event containing command execution data"""
args, unknown_args = parser_preexec.parse_known_args(cli_args)
process = self._terminal_sessions[args.pid]
event_data = {
'command': args.command,
'path': args.path,
'shell': args.shell,
'exit_code': 'unknown',
'session_id': process.unique_id
}
process.event = self._insert_event(data=event_data,
timestamp=args.timestamp)
def _precmd(self, cli_args: list) -> None:
"""Update the stored event with duration and exit_code"""
args, unknown_args = parser_precmd.parse_known_args(cli_args)
process = self._terminal_sessions[args.pid]
if process.event is None:
return
event_data = process.event.data
# Calculate time delta between preexec and precmd
timestamp = process.event.timestamp
cur_time = args.timestamp
time_delta = cur_time - timestamp
event_data['exit_code'] = args.exit_code
self._insert_event(data=event_data,
id=process.event.id,
timestamp=timestamp,
duration=time_delta)
process.event = None
def _preclose(self, args: argparse.Namespace, args_raw: list) -> None:
"""Remove pid and related data from terminal_processes_data"""
args, unknown_args = parser_preclose.parse_known_args(cli_args)
self._terminal_sessions.pop(args.pid)
def _heartbeat(self, cli_args: list) -> None:
"""Send heartbeat to activity bucket"""
args, unknown_args = parser_heartbeat.parse_known_args(cli_args)
process = self._terminal_sessions[args.pid]
event_data = {
'session_id': process.unique_id,
'shell': args.shell,
'path': args.path
}
event = Event(data=event_data, timestamp=args.timestamp)
inserted_heartbeat = self._client.heartbeat(
self._buckets['activity']['id'],
event,
pulsetime=self.pulsetime,
queued=True)
if inserted_heartbeat and inserted_heartbeat.id:
logger.debug('Successfully sent heartbeat')
def _insert_event(self, *args, **kwargs) -> Event:
"""Send event to the aw-server"""
event = Event(*args, **kwargs)
inserted_event = self._client.insert_event(
self._buckets['commands']['id'], event)
# aw-server assigned the event an id
assert inserted_event.id is not None
logger.debug("Successfully sent event")
return inserted_event
from random import random
from datetime import datetime, timedelta, timezone
from requests.exceptions import HTTPError
from aw_core.models import Event
from aw_client import ActivityWatchClient
def create_unique_event():
return Event(timestamp=datetime.now(timezone.utc),
duration=timedelta(),
data={"label": str(random())})
client = ActivityWatchClient("aw-test-client", testing=True)
client.connect()
bucket_name = "test-bucket"
bucket_etype = "test"
# Delete bucket before creating it, and handle error if it doesn't already exist
try:
client.delete_bucket(bucket_name)
except HTTPError as e:
pass
client.create_bucket(bucket_name, bucket_etype)
e1 = create_unique_event()
client.insert_event(bucket_name, e1)
print("Getting events")
events = client.get_events(bucket_name)