diff --git a/src/actions/__init__.py b/src/actions/__init__.py index 3827707173022015df58a51f57101e3804362773..f689a298030bdb793cabc394d600864289f1ba2e 100644 --- a/src/actions/__init__.py +++ b/src/actions/__init__.py @@ -1,4 +1,5 @@ from .basic_actions import * +from .docker import * from .interactive_session import * from .permissions import * from .server_management import * diff --git a/src/actions/basic_actions.py b/src/actions/basic_actions.py index fa6be86e8d8e213d2b549e29d1afc65d975e4edf..5c35909c6f12b7373d32e0bfe0bc744a5b012d6c 100644 --- a/src/actions/basic_actions.py +++ b/src/actions/basic_actions.py @@ -20,7 +20,7 @@ async def overall_help(message: Message, state=None): '* meme admin permissions - Manage servers permissions\n\n' # '* meme admin commands - Manage stored commands\n\n' '* meme admin exec {alias} {command} - Execute single command\n\n' - # '* meme admin docker - Manage docker on remote server\n\n' + '* meme admin docker - Manage docker on remote server\n\n' '* meme admin stats - See server statistics\n\n' '```', parse_mode=ParseMode.MARKDOWN @@ -33,11 +33,7 @@ async def overall_help(message: Message, state=None): ('command', r'(.*)') ]) async def bot_run_code(message: Message, alias: str, command: str, state=None): - server = await get_server_by_alias(message, alias) - if server.server_permissions.value < ServerPermissions.RCE.value \ - and str(message.from_user.id) != server.server_admin: - raise Exception('You do not have permission for remote code execution.' - ' Please ask user who added this server to change permission level.') + server = await get_server_by_alias(message, alias, ServerPermissions.RCE) results = chunks(await run_ssh_command(server, command), 4095) for res in results: await message.reply(f'```\n{res}```', parse_mode=ParseMode.MARKDOWN) diff --git a/src/actions/docker.py b/src/actions/docker.py new file mode 100644 index 0000000000000000000000000000000000000000..7dfbb576ae8fc3c6b16db69477fa15ae364dc909 --- /dev/null +++ b/src/actions/docker.py @@ -0,0 +1,133 @@ +import json + +from aiogram.types import Message, ParseMode + +from src.models.server import ServerPermissions +from src.utils import chunks +from src.utils.decorators import ( + bot_action +) +from src.utils.server_utils import get_server_by_alias +from src.utils.ssh import run_ssh_command + + +@bot_action(r'meme admin docker( help)?$') +async def overall_help(message: Message, state=None): + await message.reply( + '```\n' + '* ls {alias} - List all running containers\n\n' + '* clean {alias} - Remove exited containers\n\n' + '* stop {alias} {containers} - stop one or more containers\n\n' + '* rm {alias} {containers} - remove one or more containers\n\n' + '* up {alias} {containers} - run one or more stopped containers\n\n' + '```', + parse_mode=ParseMode.MARKDOWN + ) + + +@bot_action(r'meme admin docker ls', + params=[ + ('alias', r'[\w\d]+') + ]) +async def docker_ls(message: Message, alias, state=None): + server = await get_server_by_alias(message, alias, ServerPermissions.DOCKER) + containers_ids_list = await run_ssh_command(server, 'docker ps -aq') + if not containers_ids_list: + await message.reply('* No running containers was found *', parse_mode=ParseMode.MARKDOWN) + return + ids = ' '.join(containers_ids_list.splitlines()) + response = await run_ssh_command(server, f'docker inspect {ids}') + inspect_info = json.loads(response) + formatted_containers = "" + total = 0 + for container in inspect_info: + total += 1 + formatted_containers += ( + '```\n' + f'ID: {container["Id"][:12]}\n' + f'name: {container["Name"]}\n' + f'status: {container["State"]["Status"]}\n' + f'networks: {", ".join(container["NetworkSettings"]["Networks"].keys())}\n' + '```\n') + message_text = ('* Docker containers: *\n' + f'*Total:* {total}\n\n' + f'{formatted_containers}') + + results = chunks(message_text, 4095) + for res in results: + await message.reply(res, parse_mode=ParseMode.MARKDOWN) + + +@bot_action(r'meme admin docker clean', + params=[ + ('alias', r'[\w\d]+') + ]) +async def docker_clean(message: Message, alias, state=None): + server = await get_server_by_alias(message, alias, ServerPermissions.DOCKER) + output = await run_ssh_command(server, 'docker ps -a | grep Exit | cut -d ' ' -f 1 | xargs docker rm') + total = len(output.splitlines()) + await message.reply( + '* Removed containers: *\n\n' + f'*Total:* {total}' + '```\n' + f'{output}' + '```', + parse_mode=ParseMode.MARKDOWN + ) + + +@bot_action(r'meme admin docker rm', + params=[ + ('alias', r'[\w\d]+'), + ('containers', r'.*') + ]) +async def docker_rm(message: Message, alias, containers, state=None): + server = await get_server_by_alias(message, alias, ServerPermissions.DOCKER) + output = await run_ssh_command(server, f'docker rm {containers}') + total = len(output.splitlines()) + await message.reply( + '* Removed containers: *\n\n' + f'*Total:* {total}' + '```\n' + f'{output}' + '```', + parse_mode=ParseMode.MARKDOWN + ) + + +@bot_action(r'meme admin docker stop', + params=[ + ('alias', r'[\w\d]+'), + ('containers', r'.*') + ]) +async def docker_rm(message: Message, alias, containers, state=None): + server = await get_server_by_alias(message, alias, ServerPermissions.DOCKER) + output = await run_ssh_command(server, f'docker stop {containers}') + total = len(output.splitlines()) + await message.reply( + '* Stopped containers: *\n\n' + f'*Total:* {total}' + '```\n' + f'{output}' + '```', + parse_mode=ParseMode.MARKDOWN + ) + + +@bot_action(r'meme admin docker up', + params=[ + ('alias', r'[\w\d]+'), + ('containers', r'.*') + ]) +async def docker_rm(message: Message, alias, containers, state=None): + server = await get_server_by_alias(message, alias, ServerPermissions.DOCKER) + output = await run_ssh_command(server, f'docker start {containers}') + total = len(output.splitlines()) + await message.reply( + '* Upped containers: *\n\n' + f'*Total:* {total}' + '```\n' + f'{output}' + '```', + parse_mode=ParseMode.MARKDOWN + ) diff --git a/src/actions/interactive_session.py b/src/actions/interactive_session.py index 6f652bdfb3781547b111bfcee5954a34a7354e7f..1110a4482d34130c7d799ae8916442b792a29825 100644 --- a/src/actions/interactive_session.py +++ b/src/actions/interactive_session.py @@ -16,10 +16,11 @@ from aiogram.types import ( from src.models.crud.server_crud import fn_get_server from src.models.server import ServerPermissions from src.settings import settings +from src.utils import chunks from src.utils.debug_mode import debug_message from src.utils.decorators import bot_action from src.utils.server_utils import get_server_by_alias -from src.utils.ssh import session_manager, run_ssh_command +from src.utils.ssh import session_manager logger = logging.getLogger() @@ -108,4 +109,10 @@ async def run_interactive_command(message: Message, state: FSMContext): f'Session id: {current_state["session_id"]}\n' ) result = await session_manager.run_command(current_state['session_id'], message.text) - await message.reply(f'```\n{result if result else "Nothing to show"}```', parse_mode=ParseMode.MARKDOWN) + if not result: + await message.reply(f'```\nNothing to show```', parse_mode=ParseMode.MARKDOWN) + return + + results = chunks(result, 4095) + for res in results: + await message.reply(f'```\n{res}```', parse_mode=ParseMode.MARKDOWN) diff --git a/src/utils/server_utils.py b/src/utils/server_utils.py index 1ebb4296eca6c78eb97af44dae47c7396a8d3802..50ce0d84c7a285d5adceb389d5f98100c7837850 100644 --- a/src/utils/server_utils.py +++ b/src/utils/server_utils.py @@ -2,10 +2,13 @@ from aiogram.types import Message from src.models import Server from src.models.crud.server_crud import fn_get_server +from src.models.server import ServerPermissions from src.settings import settings -async def get_server_by_alias(message: Message, server_alias: str) -> Server: +async def get_server_by_alias(message: Message, + server_alias: str, + minimal_permission: ServerPermissions = ServerPermissions.INFO) -> Server: found_servers = await fn_get_server(settings.engine, chat_id=str(message.chat.id), server_alias=server_alias) @@ -19,4 +22,10 @@ async def get_server_by_alias(message: Message, server_alias: str) -> Server: ) else: raise Exception('No server linked to this chat.') - return found_servers[0] + server = found_servers[0] + if server.server_permissions.value < minimal_permission.value \ + and str(message.from_user.id) != server.server_admin: + raise Exception('You do not have permission for that action.\n' + f'Minimal permission level is `{minimal_permission.name}`\n' + f'Contact server administrator to change server permissions.') + return server diff --git a/src/utils/ssh.py b/src/utils/ssh.py index 3679296862ed84e49d4cfa51041c43f94e69fc3e..ff8a2c9107fc0c813b5bfeb43bc0ac63bb4ac42c 100644 --- a/src/utils/ssh.py +++ b/src/utils/ssh.py @@ -1,82 +1,42 @@ +import asyncio import logging -import re import uuid -import paramiko -from paramiko import SSHClient, ChannelFile +import asyncssh from src.models import Server logger = logging.getLogger() -async def open_ssh_session(server: Server) -> (SSHClient, ChannelFile, ChannelFile): - client = SSHClient() - client.load_system_host_keys() - client.set_missing_host_key_policy(paramiko.client.AutoAddPolicy) - client.connect( - hostname=server.server_address, - port=server.server_port, - username='root' - ) - channel = client.invoke_shell() - stdin = channel.makefile('wb') - stdout = channel.makefile('r') - - return client, stdin, stdout - - -async def run_interactive_command(command: str, client: SSHClient, std_in: ChannelFile, std_out: ChannelFile) -> str: - cmd = command.strip('\n') - std_in.write(cmd + '\n') - finish = 'EXIT STATUS: ' - echo_cmd = 'echo {} $?'.format(finish) - std_in.write(echo_cmd + '\n') - std_in.flush() - - sh_out = [] - sh_err = [] - for line in std_out: - if str(line).startswith(cmd) or str(line).startswith(echo_cmd): - # up for now filled with shell junk from stdin - shout = [] - elif str(line).startswith(finish): - # our finish command ends with the exit status - exit_status = int(str(line).rsplit(maxsplit=1)[1]) - if exit_status: - # stderr is combined with stdout. - # thus, swap sherr with shout in a case of failure. - sh_err = sh_out - sh_out = [] - break - else: - # get rid of 'coloring and formatting' special characters - sh_out.append(re.compile(r'(\x9B|\x1B\[)[0-?]*[ -/]*[@-~]').sub('', line). - replace('\b', '').replace('\r', '')) - - # first and last lines of shout/sherr contain a prompt - if sh_out and echo_cmd in sh_out[-1]: - sh_out.pop() - if sh_out and cmd in sh_out[0]: - sh_out.pop(0) - if sh_err and echo_cmd in sh_err[-1]: - sh_err.pop() - if sh_err and cmd in sh_err[0]: - sh_err.pop(0) - if sh_err: - return '\n'.join(sh_err) - return '\n'.join(sh_out) +async def open_ssh_session(server: Server) -> asyncssh.SSHClientProcess: + connection = await asyncssh.connect(server.server_address, server.server_port) + process = await connection.create_process('/bin/bash') + return process + + +async def run_interactive_command(command: str, + process: asyncssh.SSHClientProcess, + timeout=0.5) -> str: + process.stdin.write(command + '\n') + res = [] + try: + line = await asyncio.wait_for(process.stdout.readline(), timeout) + res.append(line) + while line: + logger.debug(line) + res.append(await asyncio.wait_for(process.stdout.readline(), timeout)) + except asyncio.exceptions.TimeoutError as e: + logger.exception(e) + return '\n'.join(res).strip() + return '\n'.join(res).strip() async def run_ssh_command(server: Server, command: str) -> str: - session = await open_ssh_session(server) - client = session[0] - stdin, stdout, stderr = client.exec_command(command) - err = stderr.read() - if err: - raise Exception(err.decode('utf-8')) - out = stdout.read() - return out.decode('utf-8') + process = await open_ssh_session(server) + res = await run_interactive_command(command, process) + process.close() + return res class SessionManager(object): @@ -90,10 +50,10 @@ class SessionManager(object): return rand_uuid async def run_command(self, connection_id: str, command: str): - return await run_interactive_command(command, *self.__connections[connection_id]) + return await run_interactive_command(command, self.__connections[connection_id]) def close(self, connection_id: str): - self.__connections[connection_id][0].close() + self.__connections[connection_id].close() session_manager = SessionManager()