284 lines
11 KiB
Python
Executable File
284 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
|
|
|
|
|
|
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 = sqlite3.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')
|
|
|
|
conn.commit()
|
|
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 = sqlite3.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)))
|
|
|
|
conn.commit()
|
|
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 = sqlite3.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()) |