uncia/uncia/manage.py

import json, sys

from datetime import datetime
from envbash import load_envbash
from getpass import getuser
from izzylib import DotDict, Path, boolean, logging, prompt
from izzylib.sql import Database
from os import environ as env
from urllib.parse import urlparse

from . import __version__
from .config import config, dbconfig, path, write_config
from .database import db
from .functions import fetch_actor, get_inbox
from .messages import Message


exe = f'{sys.executable} -m uncia.manage'
forbidden_keys = ['pubkey', 'privkey', 'version', 'rules']


## todo: use argparse to simplify argument parsing
class Command:
	def __init__(self, arguments):
		try:
			cmd = arguments[0]

		except IndexError:
			self.result = self.cmd_help()
			return

		args = arguments[1:] if len(arguments) > 1 else []

		try:
			self.result = self[cmd](*args)
		except InvalidCommandError:
			self.result = f'Not a valid command: {cmd}'


	def __getitem__(self, key):
		try:
			return getattr(self, f'cmd_{key}')
		except AttributeError:
			raise InvalidCommandError(f'Not a valid command: {key}')


	def cmd_help(self):
		return f'''Uncia Relay Management v{__version__}

python3 -m uncia.manage help:
    Show this message.

python3 -m uncia.manage list:
    List currently subscribed instances.

python3 -m add <actor or domain>: *
    Add an instance to the database.

python3 -m remove <actor, inbox, or domain>:
    Remove an instance and any retries associated with it from the database.

python3 -m uncia.manage config [key] [value]:
    Gets or sets the config. Specify a key and value to set a config option.
    Only specify a key to get the value of a specific option. Leave out the
    key and value to get all the options and their values.

python3 -m uncia.manage set <key> [value]:
    Set a config option to a value. If no value is specified, the default is used.

python3 -m uncia.manage request:
    List the instances currently requesting to join.

python3 -m uncia.manage accept <actor, inbox, or domain>: *
    Accept a request.

python3 -m uncia.manage deny <actor, inbox, or domain>: *
    Reject a request.

python3 -m uncia.manage rules [list]:
    List the current rules.

python3 -m uncia.manage rules add <rule>:
    Add a rule to the list.

python3 -m uncia.manage rules remove <rule number>:
    Remove a rule from the list.

python3 -m uncia.manage convert [pleroma or uncia]:
    Convert the database and config of another relay. Tries to convert a
    Pleroma Relay by default.

* = The relay needs to be running for the command to work
'''


	def cmd_config(self, key=None, value=None):
		with db.session as s:
			if key and value:
				if key in forbidden_keys:
					return f'Refusing to set "{key}"'

				if value == DefaultValue:
					value = None

				value = s.put.config(key, value)
				return f'Set {key} = "{value}"'

			elif key and not value:
				value = s.get.config(key)
				return f'Value for "{key}": {value}'

			output = 'Current config:\n'

			for key, value in s.get.config_all().items():
				if key not in forbidden_keys:
					output += f'    {key}: {value}\n'

			return output


	def cmd_set(self, key, *value):
		if not value:
			return self.cmd_config(key, DefaultValue)

		return self.cmd_config(key, ' '.join(value))


	def cmd_request(self, action=None, url=None):
		with db.session as s:
			if not action:
				instances = []

				for row in s.search('inbox'):
					if row.followid:
						instances.append(row.domain)

				if not instances:
					return 'No requests'

				text = 'Awaiting Requests:\n'
				text += '\n'.join([f'- {domain}' for domain in instances])
				return text

			else:
				instance = s.get.instance(url)

				if instance.followid:
					if action == 'accept':
						s.update(row=instance, followid=None)
						text = f'Accepted {url}'

					elif action == 'reject':
						s.remove(row=instance)
						text = f'Rejected {url}'

					else:
						return f'Not a valid request action: {action}'

					Message('accept', instance.followid, instance.actor).send(instance.inbox)
					return text


	def cmd_accept(self, url):
		return self.cmd_request('accept', url)


	def cmd_reject(self, url):
		return self.cmd_request('reject', url)


	def cmd_list(self):
		instance_list = 'Connected Instances:\n'

		with db.session as s:
				instance_list += '\n'.join([f'- {domain}' for domain in s.get.instance_list('domain')])

		return instance_list


	def cmd_add(self, actor_url):
		if not actor_url.startswith('https://'):
			actor_url = f'https://{actor_url}/actor'

		actor = fetch_actor(actor_url)
		inbox = get_inbox(actor)

		if not actor:
			return f'Failed to fetch actor at {actor_url}'

		with db.session as s:
			if s.get.instance(actor_url):
				return f'Instance already added to the relay'

			instance = s.put.instance(inbox, actor_url)

			if instance:
				return f'Added {instance.domain} to the relay'

			else:
				return f'Failed to add {actor_url} to the relay'


	def cmd_remove(self, data):
		with db.session as s:
			if s.delete.instance(data):
				return f'Instance removed: {data}'

			else:
				return f'Instance does not exist: {data}'


	def cmd_rules(self, *args):
		try:
			action = args[0]
		except IndexError:
			action = 'list'

		try:
			rule = ' '.join(args[1:])
		except IndexError:
			rule = None

		if action in ['add', 'remove'] and rule == None:
			return 'You forgot to specify a rule'

		with db.session as s:
			rules = s.get.config('rules')

			if action == 'list':
				if not rules:
					return 'No rules specified'

				text = 'Relay Rules:\n'
				text += '\n'.join([f'{idx} {rule}' for idx, rule in enumerate(rules)])

				return text

			elif action == 'add':
				rules.append(rule)
				s.put.config('rules', rules)
				return f'Added rule: {rule}'

			elif action == 'remove':
				rule = int(rule)
				rule_text = rules[rule]

				rules.remove(rule_text)
				s.put.config('rules', rules)
				return f'Removed rule: {rule_text}'


	def cmd_convert(self, relay='uncia'):
		if relay.lower() == 'uncia':
			return self.convert_uncia()

		if relay.lower() == 'pleroma':
			return self.convert_pleroma()


	def convert_uncia(self):
		cfg = load_old_uncia_config()
		pgdb = Database('postgresql', **cfg)

		new_cfg = cfg.copy()
		new_cfg.pop('name', None)
		dbconfig.update(new_cfg)

		with pgdb.session as s:
			oldcfg = DotDict({row.key: row.value for row in s.search('config')})
			inboxes = s.search('inboxes')
			requests = s.search('requests')
			retries = s.search('retries')
			users = s.search('users')
			whitelist = s.search('whitelist')
			domainbans = s.search('domainbans')
			userbans = s.search('userbans')
			actorkey = s.fetch('keys', actor='default')

		config.host = oldcfg.host
		config.listen = oldcfg.address
		config.port = oldcfg.port

		write_config()

		with db.session as s:
			s.put.config('name', oldcfg.name)
			s.put.config('admin', oldcfg.admin)
			s.put.config('show_domain_bans', boolean(oldcfg.show_domainbans))
			s.put.config('show_user_bans', boolean(oldcfg.show_userbans))
			s.put.config('whitelist', boolean(oldcfg.whitelist))
			s.put.config('block_relays', boolean(oldcfg.block_relays))
			s.put.config('require_approval', boolean(oldcfg.require_approval))
			s.put.config('privkey', actorkey.privkey)
			s.put.config('pubkey', actorkey.pubkey)

			for inbox in [*inboxes, *requests]:
				try:
					timestamp = datetime.fromtimestamp(inbox.timestamp)
				except:
					timestamp = datetime.now()

				s.insert('inbox',
					actor = inbox.actor,
					inbox = inbox.inbox,
					domain = inbox.domain,
					followid = inbox.get('followid'),
					timestamp = timestamp
				)

			for retry in retries:
				try:
					timestamp = datetime.fromtimestamp(retry.timestamp)
				except:
					timestamp = datetime.now()

				instance = s.fetch('inbox', inbox=retry.inbox)

				if not instance:
					continue

				s.insert('retry',
					inboxid = instance.id,
					msgid = retry.msgid,
					headers = json.loads(retry.headers),
					data = json.loads(retry.data),
					timestamp = timestamp
				)

			for user in users:
				try:
					timestamp = datetime.fromtimestamp(user.timestamp)
				except:
					timestamp = datetime.now()

				s.insert('user',
					handle = user.handle,
					username = user.username,
					hash = user.password,
					level = 30,
					timestamp = timestamp
				)

			for instance in whitelist:
				try:
					timestamp = datetime.fromtimestamp(instance.timestamp)
				except:
					timestamp = datetime.now()

				s.insert('whitelist',
					domain = instance.domain,
					timestamp = timestamp
				)

			for row in [*domainbans, *userbans]:
				try:
					timestamp = datetime.fromtimestamp(row.timestamp)
				except:
					timestamp = datetime.now()

				s.insert('ban',
					handle = row.get('username'),
					domain = None if row.domain.lower() == 'any' else row.domain,
					reason = row.reason,
					timestamp = timestamp
				)

		pgdb.close()

		delete_old = prompt('Config and database successfully converted. Delete old files and drop the old database?', default='no', valtype=boolean, options=['yes', 'no'])

		if delete_old:
			path.data.join('production.env').delete()

			pgdb._connect_args[0]['name'] = 'postgres'
			pgdb.open()

			with pgdb.session as s:
				s.execute(f'DROP DATABASE {cfg.name}')

			pgdb.close()

			return f'Successfully dropped database: {cfg.name}'


	def convert_pleroma(self):
		try:
			pl_cfg = load_pleroma_relay_config()
			pl_db = DotDict()
			pl_db.load_json(pl_cfg.db)

		except FileNotFoundError as e:
			return f'Error when loading config or database: {e}'

		config.listen = pl_cfg.listen
		config.port = pl_cfg.port
		config.host = pl_cfg.host

		write_config()

		for instance in pl_db['relay-list']:
			self.cmd_add(urlparse(instance).netloc)

		with db.session as s:
			s.put.config('privkey', pl_db.actorKeys.privateKey)
			s.put.config('pubkey', pl_db.actorKeys.publicKey)
			s.put.config('whitelist', pl_cfg.whitelist_enabled)
			s.put.config('description', pl_cfg.note)
			s.put.config('blocked_software', pl_cfg.blocked_software)

			for domain in pl_cfg.blocked_instances:
				s.put.ban(domain=domain)

			for domain in pl_cfg.whitelist:
				s.put.whitelist(domain)

		delete_old = prompt('Config and database successfully converted. Delete old files?', default='no', valtype=boolean, options=['yes', 'no'])

		if delete_old:
			Path('relay.yaml').delete()
			pl_cfg.db.delete()


def load_old_uncia_config():
	try:
		load_envbash(path.config.join('production.env'))

	except FileNotFoundError as e:
		pass

	cfg = DotDict(
		host = env.get('DBHOST'),
		port = int(env.get('DBPORT', 5432)),
		user = env.get('DBUSER', getuser()),
		password = env.get('DBPASS'),
		name = env.get('DBNAME', 'uncia'),
		maxconnections = int(env.get('DBCONNUM', 15))
	)

	#for k,v in cfg.items():
		#if v == None:
			#del cfg[k]

		#if k == 'host' and Path(v).exists():
			#cfg['unix_socket'] = v
			#del cfg[k]

	return cfg


def load_pleroma_relay_config():
	with open('relay.yaml') as f:
		yaml_file = yaml.load(f, Loader=yaml.FullLoader)

		if not yaml_file.get('ap'):
			raise ValueError('Missing AP section in Pleroma Relay config')

		config = DotDict({
			'db': Path(yaml_file.get('db', 'relay.jsonld')).resolve,
			'listen': yaml_file.get('listen', '127.0.0.1'),
			'port': int(yaml_file.get('port', 3621)),
			'note': yaml_file.get('note'),
			'blocked_software': [v.lower() for v in yaml_file['ap'].get('blocked_software', [])],
			'blocked_instances': yaml_file['ap'].get('blocked_instances', []),
			'host': yaml_file['ap'].get('host', 'localhost'),
			'whitelist': yaml_file['ap'].get('whitelist', []),
			'whitelist_enabled': yaml_file['ap'].get('whitelist_enabled', False)
		})

	return config


class InvalidCommandError(Exception):
	pass


class DefaultValue(object):
	pass


if __name__ == '__main__':
	args = sys.argv[1:] if len(sys.argv) > 1 else []
	cmd = Command(args)

	if cmd.result:
		print(cmd.result)