add check_matrix_msg, other minor chnages

check_federation.py

@@ -2,7 +2,6 @@
 import argparse
 import asyncio
 import json
-import os
 import sys
 import time
 import traceback
@@ -10,10 +9,10 @@ import urllib
 from datetime import datetime
 from uuid import uuid4
 
-from nio import AsyncClient, AsyncClientConfig, JoinError, JoinResponse, LoginResponse, RoomCreateError, RoomGetEventResponse, RoomSendError
+from nio import JoinError, JoinResponse, RoomCreateError, RoomGetEventResponse, RoomSendError
 
 import checker.nagios as nagios
-from checker.synapse_client import leave_all_rooms_async, leave_room_async
+from checker.synapse_client import leave_all_rooms_async, leave_room_async, login_and_cache
 
 parser = argparse.ArgumentParser(description='Test federation between two homeservers.')
 parser.add_argument('--bot1-user', required=True, help='User ID for bot 1.')
@@ -33,23 +32,6 @@ bot1_hs_domain = urllib.parse.urlparse(args.bot1_hs).netloc
 bot2_hs_domain = urllib.parse.urlparse(args.bot2_hs).netloc
 
 
-def write_details_to_disk(resp: LoginResponse, homeserver, config_file) -> None:
-    """Writes the required login details to disk so we can log in later without
-    using a password.
-    Arguments:
-        resp {LoginResponse} -- the successful client login response.
-        homeserver -- URL of homeserver, e.g. "https://matrix.example.org"
-    """
-    # open the config file in write-mode
-    with open(config_file, "w") as f:
-        # write the login details to disk
-        json.dump({"homeserver": homeserver,  # e.g. "https://matrix.example.org"
-                   "user_id": resp.user_id,  # e.g. "@user:example.org"
-                   "device_id": resp.device_id,  # device ID, 10 uppercase letters
-                   "access_token": resp.access_token,  # cryptogr. access token
-                   }, f, )
-
-
 async def test_one_direction(sender_client, receiver_client, receiver_user_id):
     # The sender creates the room and invites the receiver
     test_room_name = str(uuid4())
@@ -127,35 +109,9 @@ async def test_one_direction(sender_client, receiver_client, receiver_user_id):
     return bot1_msg_delta, nagios.OK, new_test_room_id
 
 
-async def login(user_id, passwd, homeserver, config_file=None):
-    client = AsyncClient(homeserver, user_id, config=AsyncClientConfig(request_timeout=args.timeout, max_timeout_retry_wait_time=10))
-    if config_file:
-        # If there are no previously-saved credentials, we'll use the password
-        if not os.path.exists(config_file):
-            resp = await client.login(passwd)
-
-            # check that we logged in successfully
-            if isinstance(resp, LoginResponse):
-                write_details_to_disk(resp, homeserver, config_file)
-            else:
-                print(f'UNKNOWN: failed to log in "{resp}"')
-                sys.exit(nagios.UNKNOWN)
-        else:
-            # Otherwise the config file exists, so we'll use the stored credentials
-            with open(config_file, "r") as f:
-                config = json.load(f)
-                client = AsyncClient(config["homeserver"])
-                client.access_token = config["access_token"]
-                client.user_id = config["user_id"]
-                client.device_id = config["device_id"]
-    else:
-        await client.login(passwd)
-    return client
-
-
 async def main() -> None:
-    bot1 = await login(args.bot1_user, args.bot1_pw, args.bot1_hs, args.bot1_auth_file)
+    bot1 = await login_and_cache(args.bot1_user, args.bot1_pw, args.bot1_hs, args.bot1_auth_file, request_timeout=args.timeout)
-    bot2 = await login(args.bot2_user, args.bot2_pw, args.bot2_hs, args.bot2_auth_file)
+    bot2 = await login_and_cache(args.bot2_user, args.bot2_pw, args.bot2_hs, args.bot2_auth_file, request_timeout=args.timeout)
 
     bot1_output_msg, bot1_output_code, bot1_new_room_id = await test_one_direction(bot1, bot2, args.bot2_user)
     bot2_output_msg, bot2_output_code, bot2_new_room_id = await test_one_direction(bot2, bot1, args.bot1_user)

check_matrix_msg.py

@@ -0,0 +1,90 @@
+import argparse
+import asyncio
+import sys
+import time
+
+from nio import JoinResponse, RoomSendResponse
+
+from checker import nagios, print_icinga2_check_status
+from checker.synapse_client import login_and_cache
+
+
+async def main(args):
+    start_time = time.time()
+    client = await login_and_cache(args.username, args.password, args.homeserver, auth_file=args.auth_file)
+
+    join_response = await client.join(args.room)
+    if not isinstance(join_response, JoinResponse):
+        print_icinga2_check_status(f'failed to join room!\n{join_response}', nagios.STATE_CRIT)
+        sys.exit(nagios.STATE_CRIT)
+
+    time.sleep(1)
+
+    test_msg = f'TEST MESSAGE AT {time.time()}'
+
+    send_response = await client.room_send(
+        room_id=args.room,
+        message_type="m.room.message",
+        content={"msgtype": "m.text", "body": test_msg},
+    )
+
+    if not isinstance(send_response, RoomSendResponse):
+        print_icinga2_check_status(f'failed to send message!\n{send_response}', nagios.STATE_CRIT)
+        sys.exit(nagios.STATE_CRIT)
+
+    time.sleep(1)
+
+    event_id = send_response.event_id
+
+    found_msg = False
+    for i in range(args.tries):
+        sync_response = await client.sync(30000)
+        room_timeline = sync_response.rooms.join[args.room].timeline
+        if any(event.event_id == event_id for event in room_timeline.events):
+            found_msg = True
+            break
+
+    await client.close()
+
+    if not found_msg:
+        print_icinga2_check_status(f'failed to find our message in the room timeline', nagios.STATE_CRIT)
+        sys.exit(nagios.STATE_CRIT)
+
+    elapsed_time = round((time.time() - start_time) - 2, 2)
+    perfdata = {
+        'elapsed': {
+            'value': elapsed_time,
+            'warn': args.warn,
+            'crit': args.critical,
+            'min': 0,
+            'unit': 's'
+        }
+    }
+
+    if elapsed_time >= args.critical:
+        return_code = nagios.STATE_CRIT
+        return_msg = f'message send time was {elapsed_time} seconds'
+    elif elapsed_time >= args.warn:
+        return_code = nagios.STATE_WARN
+        return_msg = f'message send time was {elapsed_time} seconds'
+    else:
+        return_code = nagios.STATE_OK
+        return_msg = 'message send successful'
+
+    print_icinga2_check_status(return_msg, return_code, perfdata=perfdata)
+    sys.exit(return_code)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Test sending message to a Matrix room.')
+    parser.add_argument('-s', '--homeserver', required=True, help='Your Matrix server. Example: https://your-homeserver.org')
+    parser.add_argument('-u', '--username', required=True, help='Matrix username.')
+    parser.add_argument('-p', '--password', required=True, help='Matrix password.')
+    parser.add_argument('-r', '--room', required=True, help='Room to send a message to.')
+    parser.add_argument('-a', '--auth-file', help="File to cache the bot's login details to.")
+    parser.add_argument('-t', '--tries', type=int, default=10, help="How many times to check for the test message.")
+    parser.add_argument('--warn', type=int, default=20, help='Warn level for message send time.')
+    parser.add_argument('--critical', type=int, default=30, help='Critical level for message send time.')
+    args = parser.parse_args()
+
+    asyncio.run(main(args))

checker/synapse_client.py

@@ -4,30 +4,15 @@ import json
 import os
 import sys
 import time
+from pathlib import Path
 
 import aiofiles.os
 import magic
 import markdown
 from PIL import Image
-from nio import AsyncClient, LoginResponse, MatrixRoom, RoomForgetResponse, RoomLeaveResponse, RoomSendError, UploadResponse
+from nio import AsyncClient, LoginResponse, MatrixRoom, RoomForgetResponse, RoomLeaveResponse, RoomSendError, UploadResponse, AsyncClientConfig
 
-from . import nagios
+from . import nagios as nagios
-
-
-def handle_err(func):
-    def wrapper(*args, **kwargs):
-        try:
-            crit, ret = func(*args, **kwargs)
-        except Exception as e:
-            print(f"UNKNOWN: exception '{e}'")
-            sys.exit(nagios.UNKNOWN)
-        if crit:
-            print(f"CRITICAL: {crit}")
-            sys.exit(nagios.CRITICAL)
-        else:
-            return ret
-
-    return wrapper
 
 
 def write_login_details_to_disk(resp: LoginResponse, homeserver, config_file) -> None:
@@ -177,3 +162,48 @@ async def leave_all_rooms_async(client, exclude_starting_with=None):
 
 def leave_all_rooms(client, exclude_starting_with=None):
     return asyncio.run(leave_all_rooms_async(client, exclude_starting_with))
+
+
+def write_details_to_disk(resp: LoginResponse, homeserver, config_file) -> None:
+    """Writes the required login details to disk so we can log in later without
+    using a password.
+    Arguments:
+        resp {LoginResponse} -- the successful client login response.
+        homeserver -- URL of homeserver, e.g. "https://matrix.example.org"
+    """
+    # open the config file in write-mode
+    with open(config_file, "w") as f:
+        # write the login details to disk
+        json.dump({"homeserver": homeserver,  # e.g. "https://matrix.example.org"
+                   "user_id": resp.user_id,  # e.g. "@user:example.org"
+                   "device_id": resp.device_id,  # device ID, 10 uppercase letters
+                   "access_token": resp.access_token,  # cryptogr. access token
+                   }, f, )
+
+
+async def login_and_cache(user_id, passwd, homeserver, auth_file=None, request_timeout=60):
+    auth_file = Path(auth_file).resolve().absolute().expanduser()
+    client = AsyncClient(homeserver, user_id, config=AsyncClientConfig(request_timeout=request_timeout, max_timeout_retry_wait_time=10))
+    if auth_file:
+        # If there are no previously-saved credentials, we'll use the password
+        if not os.path.exists(auth_file):
+            resp = await client.login(passwd)
+
+            # check that we logged in successfully
+            if isinstance(resp, LoginResponse):
+                write_details_to_disk(resp, homeserver, auth_file)
+            else:
+                # TODO: fall back to re-logging in and caching creds
+                print(f'UNKNOWN: failed to log in "{resp}"')
+                sys.exit(nagios.UNKNOWN)
+        else:
+            # Otherwise the config file exists, so we'll use the stored credentials
+            with open(auth_file, "r") as f:
+                config = json.load(f)
+                client = AsyncClient(config["homeserver"])
+                client.access_token = config["access_token"]
+                client.user_id = config["user_id"]
+                client.device_id = config["device_id"]
+    else:
+        await client.login(passwd)
+    return client