Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
import argparse
import resource
import aiohttp
import aiodns
import asyncio
from collections import Counter
import datetime
import time
import socket
import pickle
import struct
import psycopg2
async def fetch(session, url, headers=None):
# print(f'trying {url}')
async with session.get(url, ssl=False, headers=headers) as response:
return await response.json(content_type=None)
def ensure_rlimit_file():
_, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
resource.setrlimit(resource.RLIMIT_NOFILE, (hard, hard))
soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
if soft < 100000 or hard < 100000:
print(f'collection with aiohttp will not work reliably using soft {soft} and hard {hard} open file limit, increase limit for user')
async def homeserver_for_domain(session, resolver, domain):
# for spec see https://matrix.org/docs/spec/server_server/r0.1.3#resolving-server-names
# boils down to .well-known, _matrix._tcp. SRV, domain:8448
headers = None
try:
res = await fetch(session, f'https://{domain}/.well-known/matrix/server')
homeserver_initial = res['m.server']
if ':' not in homeserver_initial:
homeserver = f'{homeserver_initial}:8448'
else:
homeserver = homeserver_initial
return homeserver, headers, 'well-known'
except Exception as e:
# well-known failed, try SRV
try:
res = await resolver.query(f'_matrix._tcp.{domain}', 'SRV')
homeserver = f'{res[0].host}:{res[0].port}'
# if srv is found, request needs to include Host header for domain, not srv host. matrix.org is a good example for this
headers = {'Host': domain}
return homeserver, headers, 'SRV'
except Exception as e:
# well-known failed, SRV failed, use domain:8448 as fallback
homeserver = f'{domain}:8448'
return homeserver, headers, 'fallback'
async def version_for_homeserver(session, resolver, domain, homeserver, headers, method, debug=False):
try:
res = await fetch(session, f'https://{homeserver}/_matrix/federation/v1/version', headers)
versions = res['server']
# version number may look like 1.6.1 (abcd,branch,...), remove everything after first space with split()[0]
version_string = '{0}/{1}'.format(versions['name'], versions['version'].split()[0])
if debug:
print(f'{domain} has {homeserver} via {method} with {version_string}')
return version_string
except Exception as e:
if debug:
print(f'{domain} failed with {e} using {method}')
async def version_for_domain(session, resolver, domain, debug):
homeserver, headers, method = await homeserver_for_domain(session, resolver, domain)
return await version_for_homeserver(session, resolver, domain, homeserver, headers, method, debug)
def file_destinations(file_destinations_file='test_destinations.txt'):
with open(file_destinations_file) as f:
destinations = []
for line in f.readlines():
destinations.append(line.rstrip('\n'))
return destinations
def test_destinations():
# matrix: no well-known, srv with port, host header matrix.org necessary
# librem.one: no well-known, srv with port, host header not necessary
# synod.im: well-known with port 443
# digitale-gesellschaft.ch: well-known without port (thus implicit 8448)
# matrix.tum.de: no well-known, no srv, thus fallback
return ['matrix.org', 'synod.im', 'matrix.tum.de', 'digitale-gesellschaft.ch', 'librem.one']
def postgres_destinations():
conn = psycopg2.connect('dbname=synapse host=localhost user=matrix_stats password=matrix_stats')
cur = conn.cursor()
cur.execute('SELECT destination FROM destinations;')
destinations = []
for destination in cur:
destination = destination[0]
destinations.append(destination)
return destinations
def format_report(versions):
report = ""
total = 0
for _, count in versions.items():
total = total + count
report += '{0}\n'.format(datetime.datetime.now().isoformat())
report += f'{total} homeservers online\n\n'
for version, count in versions.items():
report += f'{count:<4} {version}\n'
return report
def graphite(versions, host='localhost', port=2004):
now = int(time.time())
tuples = ([])
for version, count in list(versions.items())[:15]:
version = version.lower().replace('.', '-').replace('/', '.')
tuples.append((version, (now, count)))
package = pickle.dumps(tuples, 2)
size = struct.pack('!L', len(package))
try:
sock = socket.socket()
sock.connect( (host, port) )
sock.sendall(size)
sock.sendall(package)
except socket.error:
print(f'Couldnt connect to {host} on port {port}, is carbon-cache.py running?')
async def main(args):
debug = args.debug
ensure_rlimit_file()
if args.enable_postgres:
destination_func = postgres_destinations
elif args.enable_file_destinations:
destination_func = file_destinations
else:
destination_func = test_destinations
timeout = aiohttp.ClientTimeout(sock_connect=5, sock_read=5)
async with aiohttp.ClientSession(timeout=timeout) as session:
resolver = aiodns.DNSResolver()
tasks = []
for domain in destination_func():
tasks.append(version_for_domain(session, resolver, domain, debug))
versions_initial = Counter(await asyncio.gather(*tasks)).most_common()
# remove "None" version count for unreachable servers
versions = {}
for version, count in versions_initial:
if version is not None:
versions[version] = count
if debug:
print('\n' + format_report(versions), end='')
if args.enable_report:
reportfile = 'reports/report-{0}.txt'.format(datetime.datetime.now(datetime.timezone.utc).isoformat(timespec='seconds'))
with open(reportfile, 'w') as f:
f.write(format_report(versions))
wwwfile = '/var/www/html/mxversions.txt'
with open(wwwfile, 'w') as f:
f.write(format_report(versions))
if args.enable_graphite:
graphite(versions)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--enable-postgres", help="use postgres destinations, if not specified, uses internal test destinations", action="store_true")
parser.add_argument("--enable-file-destinations", help="destinations from file test_destinations.txt in current folder, do not use with --enable-postgres", action="store_true")
parser.add_argument("--enable-report", help="write out a report to file", action="store_true")
parser.add_argument("--enable-graphite", help="send result to graphite on localhost:2004 via pickle", action="store_true")
parser.add_argument("--debug", help="print out report for debugging to stdout", action="store_true")
args = parser.parse_args()
loop = asyncio.get_event_loop()
loop.run_until_complete(main(args))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment