client.py

#!/usr/bin/python3

# Sequana - Onion Address Reachability Tester
#
# Copyright (C) 2025 Aaron M. D. Jones <aaron@alphachat.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.

import argparse
import asyncio
import ssl
import sys

from AlphaChat.ConfigPydle import ConfigPydleClient
from datetime import datetime, timezone
from python_socks.async_.asyncio import Proxy



class OnionServiceIncompleteResponse(Exception):
	pass

class OnionServiceIncorrectResponse(Exception):
	pass

class SequanaClient(ConfigPydleClient):

	def __init__(self, *args, **kwargs):

		super().__init__(*args, **kwargs)

		_integer_minmax = {
			'check_interval':   [  300,   86400, False ],
			'check_stagger':    [   10,     120, False ],
			'log_colour_bad':   [    0,      15,  True ],
			'log_colour_good':  [    0,      15,  True ],
			'report_interval':  [ 3600, 2419200,  True ],
		}

		_onion_service_required_keys = [
			'address',
			'port',
		]

		_onion_service_key_types = {
			'address':  [ str ],
			'port':     [ int ],
			'name':     [ str ],
			'tls':      [ bool, str ],
			'write':    [ str ],
			'read':     [ str ],
		}

		for key in _integer_minmax:
			if key not in self.acconfig:
				continue
			try:
				vmin = _integer_minmax[key][0]
				vmax = _integer_minmax[key][1]
				allz = _integer_minmax[key][2]
				var = int(self.acconfig[key])
				self.acconfig[key] = var
				if var == 0 and allz:
					continue
				if var < vmin or var > vmax:
					raise ValueError('Integer variable out of range')
			except:
				if allz:
					raise ValueError(f'Configuration item "{key}" must be an integer with a ' \
					                 f'value of either zero, or between {vmin} and {vmax} ' \
					                 f'(inclusive)')
				else:
					raise ValueError(f'Configuration item "{key}" must be an integer with a ' \
					                 f'value between {vmin} and {vmax} (inclusive)')

		if not isinstance(self.acconfig['onion_services'], list):
			raise ValueError('Configuration item "onion_services" must be a list of dicts')

		if not len(self.acconfig['onion_services']):
			raise ValueError('Configuration item "onion_services" is empty')

		for onion_service in self.acconfig['onion_services']:

			if not isinstance(onion_service, dict):
				raise ValueError('Configuration item "onion_services" must be a list of dicts')

			for key in _onion_service_required_keys:
				if key not in onion_service:
					raise ValueError(f'Configuration item "onion_services" contains an item ' \
					                 f'with missing required key "{key}"')

			for key in onion_service.keys():

				acceptable_type = False

				if key not in _onion_service_key_types:
					raise ValueError(f'Configuration item "onion_services" contains an item ' \
					                 f'with key "{key}" which is not expected')

				for type in _onion_service_key_types[key]:
					if isinstance(onion_service[key], type):
						acceptable_type = True
						break

				if not acceptable_type:
					raise ValueError(f'Configuration item "onion_services" contains an item ' \
					                 f'with key "{key}" which is not of an acceptable type')

		if not isinstance(self.acconfig['log_channel'], str):
			raise ValueError('Configuration item "log_channel" must be a string')

		self.acchannels.add(self.acconfig['log_channel'])



	async def log_message_bad(self, message):

		log_colour = self.acconfig['log_colour_bad']

		if log_colour:
			message = f'\x03{log_colour:02}{message}\x03'

		await self.message(self.acconfig['log_channel'], message)



	async def log_message_good(self, message):

		log_colour = self.acconfig['log_colour_good']

		if log_colour:
			message = f'\x03{log_colour:02}{message}\x03'

		await self.message(self.acconfig['log_channel'], message)



	async def on_autoperform_done(self):

		num = 0
		for onion_service in self.acconfig['onion_services']:
			await self.add_ev_task(self.check_onion_service_task(onion_service, num))
			num += 1



	async def check_onion_service_task(self, onion_service, num):

		report_interval = self.acconfig['report_interval']
		check_interval = self.acconfig['check_interval']
		check_stagger = self.acconfig['check_stagger']

		onion_service['total_checks'] = 0
		onion_service['total_passes'] = 0

		name = onion_service.get('name', onion_service['address'])

		while True:

			current_ts = int(datetime.now(tz=timezone.utc).timestamp())

			await asyncio.sleep(check_interval - (current_ts % check_interval))
			await asyncio.sleep(check_stagger * num)

			current_ts = int(datetime.now(tz=timezone.utc).timestamp())
			report_ts = onion_service.get('report_ts', 0)

			await self.check_onion_service_once(onion_service, num, name)

			if report_interval == 0 or (current_ts - report_ts) < report_interval:
				continue

			onion_service['report_ts'] = current_ts
			fail_count = onion_service.pop('fail_count', 0)
			pass_count = onion_service.pop('pass_count', 0)
			availability = 100 * (onion_service['total_passes'] / onion_service['total_checks'])
			availability = '{0:.2f}%'.format(availability)

			if fail_count and pass_count:
				await self.log_message_bad(f'Onion service {num}/"{name}": Heartbeat: ' \
				                           f'{pass_count} check(s) passed, ' \
				                           f'{fail_count} check(s) failed, ' \
				                           f'{availability} availability')
			elif fail_count:
				await self.log_message_bad(f'Onion service {num}/"{name}": Heartbeat: ' \
				                           f'{fail_count} check(s) failed, ' \
				                           f'{availability} availability')
			elif pass_count:
				await self.log_message_good(f'Onion service {num}/"{name}": Heartbeat: ' \
				                            f'{pass_count} check(s) passed, ' \
				                            f'{availability} availability')



	async def check_onion_service_once(self, onion_service, num, name):

		try:
			socks_host = self.acconfig['tor_address']
			socks_port = self.acconfig['tor_port']
			dest_host  = onion_service['address'] + '.onion'
			dest_port  = onion_service['port']
			dest_name  = None
			ssl_ctx    = None

			onion_service['total_checks'] += 1

			if onion_service.get('tls', None):
				ssl_ctx = ssl.create_default_context()
				if isinstance(onion_service['tls'], str):
					dest_name = onion_service['tls']
				else:
					dest_name = dest_host

			proxy          = Proxy.from_url(f'socks5://{socks_host}:{socks_port}')
			sock           = await proxy.connect(dest_host=dest_host, dest_port=dest_port)
			reader, writer = await asyncio.open_connection(host=None, port=None, sock=sock,
			                                               ssl=ssl_ctx, server_hostname=dest_name)

			if 'write' in onion_service:
				writer.write(onion_service['write'].encode('utf-8'))
				await writer.drain()

			if 'read' in onion_service:
				expected = onion_service['read'].encode('utf-8')
				try:
					response = await reader.readexactly(len(expected))
				except asyncio.IncompleteReadError as e:
					raise OnionServiceIncompleteResponse(str(e.partial))
				if response != expected:
					raise OnionServiceIncorrectResponse(str(response))

			try:
				writer.close()
			except:
				pass
			try:
				reader.close()
			except:
				pass
			try:
				await writer.wait_closed()
			except:
				pass
			try:
				await reader.wait_closed()
			except:
				pass

			onion_service['pass_count'] = onion_service.get('pass_count', 0) + 1
			onion_service['total_passes'] += 1

			if onion_service.pop('is_failing', False):
				await self.log_message_good(f'Onion service {num}/"{name}": Recovered')

		except Exception as e:
			onion_service['is_failing'] = True
			onion_service['fail_count'] = onion_service.get('fail_count', 0) + 1
			await self.log_message_bad(f'Onion service {num}/"{name}": {type(e)}: {str(e)}')



if __name__ == '__main__':

	default_config_keys = {
		'check_interval':   900,
		'check_stagger':    30,
		'log_colour_bad':   0,
		'log_colour_good':  0,
		'report_interval':  604800,
		'tor_address':      '127.0.0.1',
		'tor_port':         9050,
	}

	required_config_keys = [
		'log_channel',
		'onion_services',
	]

	parser = argparse.ArgumentParser()
	parser.add_argument('--config', default='client.yaml')
	args = parser.parse_args()

	client = SequanaClient(args.config, default_config_keys, required_config_keys)
	client.run()
	sys.exit(1)