Files
titlebot-ng-ng/get_user_list.py
T
2026-08-12 18:22:26 -05:00

285 lines
11 KiB
Python
Executable File

"""
Python script that reads configs/config.yaml and scrapes a list of active users from the Matrix room.
Stores the user list in a SQLite database named after the room.
"""
import yaml
import re
import asyncio
import sqlite3
import time
import sys
from nio import AsyncClient, LoginResponse, RoomMessagesResponse, SyncResponse
from room import Room
import dbutil
class UserListManager:
"""Manages collection and storage of active users in a Matrix room."""
def __init__(self, config_path="configs/config.yaml"):
self.room = Room(config_path)
self.users = []
self.db_path = self.get_user_db_path()
self.init_user_db()
def get_user_db_path(self):
"""Generate database path for user list, similar to suggestions database."""
safe = re.sub(r'[^a-zA-Z0-9_-]', '_', self.room.config.room_id)
return f"users_{safe}.db"
def init_user_db(self):
"""Initialize the user list database."""
conn = dbutil.connect(self.db_path)
c = conn.cursor()
c.execute('''
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id TEXT UNIQUE,
display_name TEXT,
last_seen INTEGER,
last_updated INTEGER,
num_posts INTEGER DEFAULT 0,
last_active INTEGER DEFAULT 0,
UNIQUE(user_id)
)
''')
# Add num_posts column if it doesn't exist (for existing databases)
c.execute("PRAGMA table_info(users)")
cols = [row[1] for row in c.fetchall()]
if 'num_posts' not in cols:
c.execute('ALTER TABLE users ADD COLUMN num_posts INTEGER DEFAULT 0')
if 'last_active' not in cols:
c.execute('ALTER TABLE users ADD COLUMN last_active INTEGER DEFAULT 0')
# Table to track when user list was last fetched
c.execute('''
CREATE TABLE IF NOT EXISTS fetch_log (
id INTEGER PRIMARY KEY AUTOINCREMENT,
timestamp INTEGER,
user_count INTEGER,
datetime TEXT,
most_recent_post INTEGER DEFAULT 0
)
''')
# Add most_recent_post column if it doesn't exist
c.execute("PRAGMA table_info(fetch_log)")
cols = [row[1] for row in c.fetchall()]
if 'most_recent_post' not in cols:
c.execute('ALTER TABLE fetch_log ADD COLUMN most_recent_post INTEGER DEFAULT 0')
dbutil.commit(conn)
conn.close()
async def get_room_members(self):
"""Fetch all members from the Matrix room and count their posts."""
# Ensure we're logged in
if not await self.room.login_to_matrix():
print("Failed to login to Matrix")
return []
# Sync to get room state
sync_resp = await self.room.client.sync(timeout=3000)
if not isinstance(sync_resp, SyncResponse):
print(f"Sync failed: {sync_resp}")
# Get the room object
room_obj = self.room.client.rooms.get(self.room.config.room_id)
if not room_obj:
print(f"Could not find room {self.room.config.room_id}")
return []
# Count posts by fetching message history
print("Counting posts per user (this may take a moment)...")
post_counts, last_activity, most_recent_post = await self.count_user_posts()
users = []
for user_id in room_obj.users:
display_name = await self.room.get_display_name(user_id)
users.append({
'user_id': user_id,
'display_name': display_name,
'last_seen': int(time.time()),
'num_posts': post_counts.get(user_id, 0),
'last_active': last_activity.get(user_id, 0)
})
self.users = users
self.most_recent_post = most_recent_post
return users
async def count_user_posts(self):
"""Count total number of posts by each user in the room and track last activity."""
post_counts = {}
last_activity = {}
most_recent_post = 0
batch_size = self.room.config.batch_size
lookback_limit = self.room.config.user_lookback_limit
next_batch = None
fetched = 0
# Fetch messages in batches to count posts per user
while fetched < lookback_limit:
try:
response = await self.room.client.room_messages(
room_id=self.room.config.room_id,
start=next_batch,
limit=batch_size
)
if not isinstance(response, RoomMessagesResponse) or not response.chunk:
break
fetched += len(response.chunk)
for event in response.chunk:
if hasattr(event, 'sender') and hasattr(event, 'body'):
sender = event.sender
post_counts[sender] = post_counts.get(sender, 0) + 1
# Track the most recent message timestamp for each user
if hasattr(event, 'server_timestamp'):
event_time = int(event.server_timestamp // 1000)
if sender not in last_activity or event_time > last_activity[sender]:
last_activity[sender] = event_time
# Track the absolute most recent post in the room
if event_time > most_recent_post:
most_recent_post = event_time
next_batch = response.end
# Stop if we've reached the end or the lookback limit
if not next_batch or fetched >= lookback_limit:
break
except Exception as e:
print(f"Error fetching messages: {e}")
break
return post_counts, last_activity, most_recent_post
def write_users_to_db(self):
"""Write collected users to the database."""
if not self.users:
print("No users to write to database")
return
conn = dbutil.connect(self.db_path)
c = conn.cursor()
current_time = int(time.time())
for user in self.users:
try:
c.execute('''
INSERT INTO users (user_id, display_name, last_seen, last_updated, num_posts, last_active)
VALUES (?, ?, ?, ?, ?, ?)
ON CONFLICT(user_id) DO UPDATE SET
display_name = excluded.display_name,
last_seen = excluded.last_seen,
last_updated = excluded.last_updated,
num_posts = excluded.num_posts,
last_active = excluded.last_active
''', (
user['user_id'],
user['display_name'],
user['last_seen'],
current_time,
user['num_posts'],
user.get('last_active', 0)
))
except sqlite3.IntegrityError:
# Update existing record
c.execute('''
UPDATE users
SET display_name = ?, last_seen = ?, last_updated = ?, num_posts = ?, last_active = ?
WHERE user_id = ?
''', (
user['display_name'],
user['last_seen'],
current_time,
user['num_posts'],
user.get('last_active', 0),
user['user_id']
))
# Log this fetch
from datetime import datetime
c.execute('''
INSERT INTO fetch_log (timestamp, user_count, datetime, most_recent_post)
VALUES (?, ?, ?, ?)
''', (current_time, len(self.users), datetime.now().isoformat(), getattr(self, 'most_recent_post', 0)))
dbutil.commit(conn)
conn.close()
print(f"Wrote {len(self.users)} users to {self.db_path}")
def read_users_from_db(self):
"""Read all users from the database."""
conn = dbutil.connect(self.db_path)
c = conn.cursor()
c.execute('SELECT user_id, display_name, last_seen, last_updated, num_posts, last_active FROM users ORDER BY num_posts DESC, display_name')
users = []
for row in c.fetchall():
users.append({
'user_id': row[0],
'display_name': row[1],
'last_seen': row[2],
'last_updated': row[3],
'num_posts': row[4],
'last_active': row[5] if len(row) > 5 else 0
})
conn.close()
return users
async def close(self):
"""Close Matrix client connection."""
await self.room.close()
async def main():
"""Main function to fetch and store user list continuously."""
config_path = sys.argv[1] if len(sys.argv) > 1 else "configs/config.yaml"
manager = UserListManager(config_path)
try:
print("Starting user list collection service...")
print("This will run continuously, updating the user list every 5 minutes.")
while True:
try:
print("\nFetching user list from Matrix room...")
users = await manager.get_room_members()
if users:
print(f"Found {len(users)} users in room")
manager.write_users_to_db()
# Display top users by post count
print("\nTop users by post count:")
sorted_users = sorted(users, key=lambda u: u['num_posts'], reverse=True)
for user in sorted_users[:10]: # Show top 10
print(f" - {user['display_name']}: {user['num_posts']} posts")
if len(sorted_users) > 10:
print(f" ... and {len(sorted_users) - 10} more users")
else:
print("No users found or could not connect to room")
# Wait 5 minutes before next update
print("\nWaiting 5 minutes before next update...")
await asyncio.sleep(300) # 5 minutes
except Exception as e:
print(f"Error during user list update: {e}")
print("Retrying in 1 minute...")
await asyncio.sleep(60) # Wait 1 minute on error before retry
except KeyboardInterrupt:
print("\nShutting down user list service...")
finally:
await manager.close()
if __name__ == "__main__":
asyncio.run(main())