mypointvpn/src/App/tests_bot.py

import datetime
from unittest import mock

from django.test import TestCase
from django.utils import timezone
from dramatiq.middleware.time_limit import TimeLimitExceeded
from Telegram import interfaces
from Telegram.models import Bot, Chat, User as TelegramUser
from Telegram.testing import Conversation, FakeApi

from . import fields, services, tasks
from .bot import texts
from .bot.interface import MyPointInterface, is_proxy_url, parse_credentials
from .models import Credentials, Point, Server
from .tests_services import FakeRunner, FakeXuiApi, create_point, patch_runner, patch_xui_api


class ParsersTests(TestCase):
	def test_credentials(self):
		self.assertEqual(parse_credentials('1.2.3.4:pa:ss').password, 'pa:ss')
		self.assertIsNone(parse_credentials('1.2.3.4:'))
		self.assertIsNone(parse_credentials('host:pw'))
		self.assertIsNone(parse_credentials('1.2.3.400:pw'))

	def test_proxy_url(self):
		self.assertTrue(is_proxy_url('http://127.0.0.1:9050'))
		self.assertTrue(is_proxy_url('socks5://user:pw@proxy.example.com:1080'))
		self.assertFalse(is_proxy_url('http://127.0.0.1'))
		self.assertFalse(is_proxy_url('ftp://127.0.0.1:21'))


class BotTestCase(TestCase):
	def setUp(self):
		self.bot = Bot.objects.create(token='1:test', interface=interfaces.key(MyPointInterface))
		self.chat = Conversation(self.bot)
		self.api = self.chat.api

	def say(self, text):
		return self.chat.say(text)

	@property
	def keyboard(self):
		return self.api.last_keyboard


class MainMenuTests(BotTestCase):
	def test_start_and_main_menu_button_are_same(self):
		self.assertEqual(self.say('/start'), texts.GREETING)
		self.assertEqual(self.keyboard, texts.MAIN_KEYBOARD)
		self.assertEqual(self.say(texts.MAIN_MENU), texts.GREETING)

	def test_lists_points(self):
		self.say('/start')
		point = create_point(TelegramUser.objects.get())
		self.assertIn(point.name, self.say('/start'))
		self.assertEqual(self.keyboard, texts.MAIN_KEYBOARD_WITH_POINTS)

	def test_faq_and_help(self):
		self.assertEqual(self.say(texts.FAQ), texts.FAQ_TEXT)
		self.assertEqual(self.say(texts.HELP), texts.HELP_TEXT)


@mock.patch.object(tasks.install_point, 'send')
class AddServerTests(BotTestCase):
	def fill(self, role_button, credentials):
		self.say(role_button)
		return self.say(credentials)

	def test_full_flow(self, send):
		self.say(texts.ADD_SERVER)
		self.assertEqual(self.keyboard, texts.DRAFT_KEYBOARD)
		card = self.fill(texts.SERVER, '10.0.0.1:panel-pw')
		self.assertIn('Сервер: 10.0.0.1', card)
		self.assertNotIn('panel-pw', card)
		self.fill(texts.MASK, '10.0.0.2:mask-pw')
		self.say(texts.DOMAIN)
		self.say('Example.COM.')
		self.say(texts.PROXY)
		card = self.say(texts.CF_WARP)
		self.assertIn('Домен: example.com', card)
		self.assertIn('Прокси: cf_warp', card)

		self.assertEqual(self.say(texts.INSTALL), texts.INSTALL_STARTED)
		point = Point.objects.get()
		self.assertEqual((point.panel.address, point.panel.root_password), ('10.0.0.1', 'panel-pw'))
		self.assertEqual((point.mask.address, point.mask.root_password), ('10.0.0.2', 'mask-pw'))
		self.assertEqual((point.domain, point.outbound_proxy), ('example.com', 'cf_warp'))
		self.assertEqual(point.owner, TelegramUser.objects.get())
		send.assert_called_once_with(point.pk, mock.ANY)
		point.refresh_from_db()
		self.assertEqual((point.task_chat, point.task_token), (Chat.objects.get(), send.call_args.args[1]))
		# сообщения с паролями удалены из чата
		self.assertEqual(len(self.api.deleted), 2)
		self.assertEqual(Chat.objects.get().state, '')

	def test_validation(self, send):
		self.say(texts.ADD_SERVER)
		self.assertEqual(self.say(texts.INSTALL), texts.DRAFT_INCOMPLETE)
		self.assertEqual(self.fill(texts.SERVER, 'nonsense'), texts.CREDENTIALS_INVALID)
		self.say('10.0.0.1:pw')
		self.assertEqual(self.fill(texts.MASK, '10.0.0.1:pw'), texts.ADDRESS_DUPLICATE)
		create_point(TelegramUser.objects.get())
		self.assertEqual(self.say('10.0.0.2:pw'), texts.ADDRESS_TAKEN.format(address='10.0.0.2'))
		self.say(texts.CANCEL)
		self.assertEqual(self.fill(texts.DOMAIN, 'not a domain'), texts.DOMAIN_INVALID)
		self.say(texts.CANCEL)
		self.assertEqual(self.fill(texts.PROXY, 'http://nope'), texts.PROXY_INVALID)
		send.assert_not_called()

	def test_install_without_mask(self, send):
		self.say(texts.ADD_SERVER)
		self.fill(texts.SERVER, '10.0.0.1:pw')
		self.assertEqual(self.say(texts.INSTALL), texts.INSTALL_STARTED)
		point = Point.objects.get()
		self.assertIsNone(point.mask)
		send.assert_called_once_with(point.pk, mock.ANY)

	def test_mask_requires_domain(self, send):
		self.say(texts.ADD_SERVER)
		self.fill(texts.SERVER, '10.0.0.1:pw')
		self.fill(texts.MASK, '10.0.0.2:pw')
		self.assertEqual(self.say(texts.INSTALL), texts.MASK_WITHOUT_DOMAIN)
		self.say(texts.MASK)
		self.assertIn('Сервер маскировки: -', self.say(texts.DELETE))
		self.assertEqual(self.say(texts.INSTALL), texts.INSTALL_STARTED)
		self.assertIsNone(Point.objects.get().mask)

	def test_cancel_in_prompt_returns_to_draft(self, send):
		self.say(texts.ADD_SERVER)
		self.fill(texts.SERVER, '10.0.0.1:pw')
		self.say(texts.MASK)
		self.assertIn('Сервер: 10.0.0.1', self.say(texts.CANCEL))
		self.assertEqual(self.say(texts.CANCEL), texts.GREETING)


class PointMenuTests(BotTestCase):
	def setUp(self):
		super().setUp()
		self.say('/start')
		self.point = create_point(TelegramUser.objects.get())
		self.point.set_status(Point.Status.READY)

	def open_point(self):
		self.say(texts.MANAGE)
		return self.say(self.point.name)

	def test_select_point(self):
		self.assertIn(self.point.name, self.say(texts.MANAGE))
		self.assertEqual(self.say('unknown'), texts.POINT_NOT_FOUND)
		card = self.say(self.point.name)
		self.assertIn('Состояние: активен', card)
		self.assertEqual(self.keyboard, texts.POINT_KEYBOARD)

	def test_foreign_point_is_not_found(self):
		other = Conversation(self.bot, user_id=8)
		other.say(texts.MANAGE)
		self.assertEqual(other.say(self.point.name), texts.POINT_NOT_FOUND)

	@mock.patch.object(tasks.release_point, 'send')
	def test_release(self, send):
		self.open_point()
		self.say(texts.RELEASE)
		self.assertIn(self.point.name, self.say(texts.NO))
		self.assertIn('Продолжить?', self.say(texts.RELEASE))
		self.assertEqual(self.say(texts.YES), texts.RELEASE_STARTED)
		send.assert_called_once_with(self.point.pk, mock.ANY)
		self.point.refresh_from_db()
		self.assertEqual(self.point.status, Point.Status.RELEASING)

	@mock.patch.object(tasks.release_point, 'send')
	def test_busy_point(self, send):
		self.point.set_status(Point.Status.INSTALLING)
		self.open_point()
		self.assertEqual(self.say(texts.RELEASE), texts.POINT_BUSY.format(status='устанавливается'))
		send.assert_not_called()

	@mock.patch.object(tasks.reinstall_point, 'send')
	def test_reinstall_one_server(self, send):
		self.open_point()
		self.say(texts.REINSTALL)
		self.assertEqual(self.say('10.9.9.9:pw'), texts.REINSTALL_UNKNOWN_ADDRESS.format(address='10.9.9.9'))
		self.assertEqual(self.say('10.0.0.2:new-pw'), texts.REINSTALL_STARTED)
		send.assert_called_once_with(self.point.pk, mock.ANY, '10.0.0.2', 'new-pw')
		self.assertEqual(len(self.api.deleted), 2)

	@mock.patch.object(tasks.reinstall_point, 'send')
	def test_reinstall_with_old_passwords(self, send):
		self.open_point()
		self.say(texts.REINSTALL)
		self.say(texts.USE_OLD)
		send.assert_called_once_with(self.point.pk, mock.ANY, None, None)

	def test_server_menu_button_opens_latest_point(self):
		self.assertIn(self.point.name, self.say(texts.SERVER_MENU))


class KeysMenuTests(BotTestCase):
	def setUp(self):
		super().setUp()
		self.say('/start')
		self.point = create_point(TelegramUser.objects.get(), mask=None)
		self.point.panel_inbound_id = 1
		self.point.status = Point.Status.READY
		self.point.save()
		self.panel_api = FakeXuiApi()
		patcher = patch_xui_api(self.panel_api)
		patcher.start()
		self.addCleanup(patcher.stop)
		self.say(texts.MANAGE)
		self.say(self.point.name)

	def test_issue_show_and_revoke_key(self):
		self.assertIn('Ключей пока нет', self.say(texts.CLIENTS))
		self.assertEqual(self.keyboard, texts.KEYS_KEYBOARD)
		self.assertEqual(self.say(texts.ADD_KEY), texts.KEY_NAME_PROMPT)
		self.assertIn('латиница', self.say('Мама'))
		card = self.say('Mom-Phone')
		self.assertIn(self.point.subscription_url('sub-mom-phone'), card)
		self.assertIn('mom-phone', self.say(texts.CLIENTS))
		self.assertIn('↓ 3.0 МБ', self.say('mom-phone'))
		self.assertEqual(self.say('ghost'), texts.KEY_NOT_FOUND)
		self.say(texts.REVOKE_KEY)
		self.assertIn('Ключей пока нет', self.say('mom-phone'))
		self.assertEqual(self.panel_api.clients, {})

	def test_panel_error(self):
		self.say(texts.CLIENTS)
		self.say(texts.ADD_KEY)
		self.panel_api.down = True
		self.assertEqual(self.say('mom'), texts.PANEL_ERROR.format(error='connection refused'))

	def test_not_ready_point(self):
		self.point.set_status(Point.Status.ERROR)
		self.assertEqual(self.say(texts.CLIENTS), texts.POINT_NOT_READY)

	def test_card_counts_keys(self):
		self.say(texts.CLIENTS)
		self.say(texts.ADD_KEY)
		self.say('mom')
		self.assertIn('Ключи: 1', self.say(texts.SERVER_MENU))


class PointSettingsTests(BotTestCase):
	def setUp(self):
		super().setUp()
		self.say('/start')
		self.point = create_point(TelegramUser.objects.get(), mask=None)
		self.point.panel_inbound_id = 1
		self.point.status = Point.Status.READY
		self.point.save()
		self.say(texts.MANAGE)
		self.say(self.point.name)

	@mock.patch.object(tasks.set_domain, 'send')
	def test_change_domain(self, send):
		self.assertIn('10.0.0.1', self.say(texts.POINT_DOMAIN))
		self.assertEqual(self.keyboard, [[texts.CANCEL, texts.DELETE]])
		self.assertEqual(self.say('not a domain'), texts.DOMAIN_INVALID)
		self.assertEqual(self.say('VPN.Example.com.'), texts.CHANGE_STARTED)
		send.assert_called_once_with(self.point.pk, mock.ANY, 'vpn.example.com')
		self.point.refresh_from_db()
		self.assertEqual(self.point.status, Point.Status.UPDATING)
		self.assertEqual(self.say(texts.SERVER_MENU).count('настраивается'), 1)
		self.assertEqual(self.say(texts.POINT_DOMAIN), texts.POINT_BUSY.format(status='настраивается'))

	@mock.patch.object(tasks.set_domain, 'send')
	def test_mask_point_cannot_drop_domain(self, send):
		Server.objects.create_server(self.point, Server.Role.MASK, Credentials('10.0.0.3', 'pw'))
		self.point.domain = 'vpn.example.com'
		self.point.save()
		self.assertIn('10.0.0.3', self.say(texts.POINT_DOMAIN))
		self.assertEqual(self.say(texts.DELETE), texts.DOMAIN_REQUIRED_BY_MASK)
		send.assert_not_called()

	@mock.patch.object(tasks.set_domain, 'send')
	def test_drop_domain(self, send):
		self.point.domain = 'vpn.example.com'
		self.point.save()
		self.say(texts.POINT_DOMAIN)
		self.say(texts.DELETE)
		send.assert_called_once_with(self.point.pk, mock.ANY, '')

	@mock.patch.object(tasks.set_outbound_proxy, 'send')
	def test_change_proxy(self, send):
		self.say(texts.POINT_PROXY)
		self.assertEqual(self.keyboard, [[texts.CF_WARP, texts.CANCEL, texts.DELETE]])
		self.assertEqual(self.say('ftp://x:21'), texts.PROXY_INVALID)
		self.assertEqual(self.say('socks5://user:secret@proxy.example:1080'), texts.CHANGE_STARTED)
		# сообщение с паролем удалено, в брокер прокси уходит зашифрованным
		self.assertEqual(len(self.api.deleted), 2)
		encrypted = send.call_args.args[2]
		self.assertNotIn('secret', encrypted)
		self.assertEqual(fields.decrypt(encrypted), 'socks5://user:secret@proxy.example:1080')

	@mock.patch.object(tasks.set_outbound_proxy, 'send')
	def test_warp_and_direct(self, send):
		self.say(texts.POINT_PROXY)
		self.say(texts.CF_WARP)
		self.assertEqual(fields.decrypt(send.call_args.args[2]), Point.CF_WARP)
		self.point.set_status(Point.Status.READY)
		self.say(texts.SERVER_MENU)
		self.say(texts.POINT_PROXY)
		self.say(texts.DELETE)
		self.assertEqual(send.call_args.args[2], '')

	def test_cancel_returns_to_point(self):
		self.say(texts.POINT_PROXY)
		self.assertIn('Состояние: активен', self.say(texts.CANCEL))
		self.assertEqual(self.keyboard, texts.POINT_KEYBOARD)

	def test_not_ready_point(self):
		self.point.set_status(Point.Status.ERROR)
		self.assertEqual(self.say(texts.POINT_PROXY), texts.POINT_NOT_READY)

	def test_card_hides_proxy_password(self):
		self.point.outbound_proxy = 'socks5://user:secret@proxy.example:1080'
		self.point.save()
		card = self.say(texts.SERVER_MENU)
		self.assertIn('socks5://***@proxy.example:1080', card)
		self.assertNotIn('secret', card)
		self.assertNotIn('user', card)


class TasksTests(TestCase):
	def setUp(self):
		bot = Bot.objects.create(token='1:test')
		owner = TelegramUser.objects.create(tg_id=1)
		self.chat = Chat.objects.create(bot=bot, chat_id=1, type='private', user=owner)
		self.point = create_point(owner)
		self.api = FakeApi()

	def run_task(self, actor, *args, status=Point.Status.INSTALLING, runner=None):
		"""Запускает задачу, как бот, и выполняет ее сообщение."""
		self.point.refresh_from_db()
		with mock.patch.object(actor, 'send') as send:
			tasks.start(actor, self.point, status, self.chat, *args)
		with (
			patch_runner(runner or FakeRunner()),
			mock.patch.object(Bot, 'api', return_value=self.api),
			mock.patch.object(tasks.check_point, 'send'),
		):
			actor.fn(*send.call_args.args)
		return self.api.last_text

	def test_install_notifies_success(self):
		text = self.run_task(tasks.install_point)
		self.assertIn(self.point.name, text)
		self.assertIn('10.0.0.2', text)
		self.assertIn('готов к работе', text)
		self.assertEqual(self.api.last_keyboard, [[texts.SERVER_MENU]])

	def test_install_notifies_failure(self):
		with self.assertLogs('App.tasks', 'ERROR'):
			text = self.run_task(tasks.install_point, runner=FakeRunner(fail='xui'))
		self.assertIn('Не удалось', text)
		self.assertIn('boom', text)

	def test_release_notifies_passwords(self):
		text = self.run_task(tasks.release_point, status=Point.Status.RELEASING)
		self.assertIn('10.0.0.1 - пароль root', text)
		self.assertFalse(Point.objects.exists())

	def test_change_notifies_result(self):
		self.point.panel_inbound_id = 1
		self.point.status = Point.Status.READY
		self.point.save()
		proxy = fields.encrypt('socks5://user:secret@proxy.example:1080')
		text = self.run_task(tasks.set_outbound_proxy, proxy, status=Point.Status.UPDATING)
		self.assertIn('socks5://***@proxy.example', text)
		with self.assertLogs('App.tasks', 'ERROR'):
			text = self.run_task(
				tasks.set_domain, 'vpn.example.com', status=Point.Status.UPDATING, runner=FakeRunner(fail='xui_cert'),
			)
		self.assertIn('прежними настройками', text)
		self.point.refresh_from_db()
		self.assertEqual((self.point.status, self.point.domain), (Point.Status.READY, ''))


class SupervisorTests(TestCase):
	"""Задача, которая прервалась, не оставляет точку занятой навсегда."""

	def setUp(self):
		bot = Bot.objects.create(token='1:test')
		owner = TelegramUser.objects.create(tg_id=1)
		self.chat = Chat.objects.create(bot=bot, chat_id=1, type='private', user=owner)
		self.point = create_point(owner)
		self.api = FakeApi()
		patcher = mock.patch.object(Bot, 'api', return_value=self.api)
		patcher.start()
		self.addCleanup(patcher.stop)

	def start(self, actor=tasks.install_point, status=Point.Status.INSTALLING, *args):
		with mock.patch.object(actor, 'send') as send, mock.patch.object(tasks.supervise, 'send_with_options') as sup:
			tasks.start(actor, self.point, status, self.chat, *args)
		sup.assert_called_once_with(args=('App.Point', self.point.pk, send.call_args.args[1]), delay=tasks.SUPERVISE_DELAY_MS)
		return send.call_args.args

	def supervise(self, token):
		with mock.patch.object(tasks.supervise, 'send_with_options') as again:
			tasks.supervise.fn('App.Point', self.point.pk, token)
		self.point.refresh_from_db()
		return again.called

	def age_heartbeat(self, delta):
		Point.objects.filter(pk=self.point.pk).update(task_heartbeat=timezone.now() - delta)

	def test_stuck_task_is_aborted(self):
		_, token = self.start()
		Point.objects.claim_task(self.point.pk, token)
		self.assertTrue(self.supervise(token))
		self.age_heartbeat(tasks.TASK_STALE + datetime.timedelta(seconds=1))
		self.assertFalse(self.supervise(token))
		self.assertEqual(self.point.status, Point.Status.ERROR)
		self.assertIn('Переустановить', self.point.error)
		self.assertEqual(self.point.task_token, '')
		self.assertIn('установка прервалась', self.api.last_text)

	def test_queued_task_waits_longer(self):
		_, token = self.start()
		self.age_heartbeat(tasks.TASK_STALE * 2)
		self.assertTrue(self.supervise(token))
		self.age_heartbeat(tasks.QUEUE_TIMEOUT + datetime.timedelta(seconds=1))
		self.assertFalse(self.supervise(token))
		self.assertEqual(self.point.status, Point.Status.ERROR)
		# сообщение все же дошло: задача снята, работы нет
		runner = FakeRunner()
		with patch_runner(runner):
			tasks.install_point.fn(self.point.pk, token)
		self.assertEqual(runner.calls, [])

	def test_interrupted_update_keeps_point_ready(self):
		token = self.start(tasks.set_domain, Point.Status.UPDATING, 'vpn.example.com')[1]
		Point.objects.claim_task(self.point.pk, token)
		self.age_heartbeat(tasks.TASK_STALE * 2)
		self.supervise(token)
		self.assertEqual((self.point.status, self.point.domain), (Point.Status.READY, ''))
		self.assertIn('прервалась', self.point.error)

	def test_supervisor_stops_after_task(self):
		args = self.start()
		with patch_runner(FakeRunner()), mock.patch.object(tasks.check_point, 'send'):
			tasks.install_point.fn(*args)
		self.assertFalse(self.supervise(args[1]))
		self.assertEqual((self.point.status, self.point.task_token), (Point.Status.READY, ''))

	def test_redelivered_message_is_not_repeated(self):
		# воркер упал посреди установки, RabbitMQ доставил сообщение снова
		args = self.start()
		Point.objects.claim_task(*args)
		runner = FakeRunner()
		with patch_runner(runner):
			tasks.install_point.fn(*args)
		self.assertEqual(runner.calls, [])
		self.point.refresh_from_db()
		self.assertEqual(self.point.status, Point.Status.ERROR)

	def test_time_limit_aborts_task(self):
		args = self.start()
		with (
			mock.patch.object(services, 'install_point', side_effect=TimeLimitExceeded),
			self.assertRaises(TimeLimitExceeded),
			self.assertLogs('App.tasks', 'ERROR'),
		):
			tasks.install_point.fn(*args)
		self.point.refresh_from_db()
		self.assertEqual((self.point.status, self.point.task_token), (Point.Status.ERROR, ''))

	def test_failed_notification_keeps_result(self):
		args = self.start()
		with (
			patch_runner(FakeRunner()),
			mock.patch.object(tasks, 'schedule_checks', side_effect=RuntimeError('boom')),
			self.assertRaises(RuntimeError),
			self.assertLogs('App.tasks', 'ERROR'),
		):
			tasks.install_point.fn(*args)
		self.point.refresh_from_db()
		self.assertEqual((self.point.status, self.point.task_token), (Point.Status.READY, ''))


class CheckTests(TestCase):
	def setUp(self):
		bot = Bot.objects.create(token='1:test')
		owner = TelegramUser.objects.create(tg_id=1)
		Chat.objects.create(bot=bot, chat_id=1, type='private', user=owner)
		self.point = create_point(owner)
		self.point.status = Point.Status.READY
		self.point.save()
		self.api = FakeApi()
		patcher = mock.patch.object(Bot, 'api', return_value=self.api)
		patcher.start()
		self.addCleanup(patcher.stop)
		with mock.patch.object(tasks.check_point, 'send') as send:
			tasks.schedule_checks(self.point)
		self.token = send.call_args.args[1]

	def check(self, problem='', token=None):
		with (
			mock.patch.object(services, 'check_point', return_value=problem) as check,
			mock.patch.object(tasks.check_point, 'send_with_options') as again,
		):
			tasks.check_point.fn(self.point.pk, token or self.token)
		self.point.refresh_from_db()
		return check.called, again.called

	def test_checks_once_a_day_and_reports_changes(self):
		self.assertEqual(self.check('панель не отвечает'), (True, True))
		self.assertEqual(self.point.health_error, 'панель не отвечает')
		self.assertIn('панель не отвечает', self.api.last_text)
		self.assertIn('Проверка: панель не отвечает', texts.point_card(self.point))
		# до следующих суток не проверяет, но цепочка живет
		self.assertEqual(self.check(), (False, True))
		Point.objects.filter(pk=self.point.pk).update(checked_at=timezone.now() - tasks.CHECK_INTERVAL)
		sent = len(self.api.sent)
		self.check('панель не отвечает')
		self.assertEqual(len(self.api.sent), sent)
		Point.objects.filter(pk=self.point.pk).update(checked_at=timezone.now() - tasks.CHECK_INTERVAL)
		self.check()
		self.assertIn('снова все в порядке', self.api.last_text)

	def test_busy_point_is_not_checked(self):
		self.point.set_status(Point.Status.INSTALLING)
		self.assertEqual(self.check(), (False, True))

	def test_old_chain_stops(self):
		with mock.patch.object(tasks.check_point, 'send'):
			tasks.schedule_checks(self.point)
		self.assertEqual(self.check(token=self.token), (False, False))
		point_id = self.point.pk
		self.point.delete()
		with mock.patch.object(tasks.check_point, 'send_with_options') as again:
			tasks.check_point.fn(point_id, self.token)
		again.assert_not_called()