diff --git a/src/actions/__init__.py b/src/actions/__init__.py index 23a65a80c4baa198215f907419bd3838e80d10ad..4627ed13cecb69fab73a79d42a634c3d014c90ec 100644 --- a/src/actions/__init__.py +++ b/src/actions/__init__.py @@ -5,6 +5,7 @@ from .permissions import * from .server_management import * from .stats import * from .systemctl import * +from .scp import * from ..models.crud.server_crud import fn_get_all_chat_ids diff --git a/src/actions/basic_actions.py b/src/actions/basic_actions.py index 9b48d46518407851ad96c16c6fa2772e51aa68a6..dcec527dba9e87ebb7057490a18121e62ca4da5b 100644 --- a/src/actions/basic_actions.py +++ b/src/actions/basic_actions.py @@ -5,7 +5,8 @@ from src.models.server import ServerPermissions from src.settings import settings from src.utils import chunks from src.utils.decorators import ( - bot_action + bot_action, + BotParam ) from src.utils.server_utils import get_server_by_alias from src.utils.ssh import run_ssh_command @@ -30,8 +31,8 @@ async def overall_help(message: Message, state=None): @bot_action(trigger_str='meme admin exec', params=[ - ('alias', r'[\w\d]+'), - ('command', r'(.*)') + BotParam('alias', r'[\w\d]+'), + BotParam('command', r'(.*)') ]) async def bot_run_code(message: Message, alias: str, command: str, state=None): server = await get_server_by_alias(message, alias, ServerPermissions.RCE) diff --git a/src/actions/docker.py b/src/actions/docker.py index 7dfbb576ae8fc3c6b16db69477fa15ae364dc909..d99fc49b6ea4d7357f5a1b1e2305533e897eb240 100644 --- a/src/actions/docker.py +++ b/src/actions/docker.py @@ -5,7 +5,8 @@ from aiogram.types import Message, ParseMode from src.models.server import ServerPermissions from src.utils import chunks from src.utils.decorators import ( - bot_action + bot_action, + BotParam ) from src.utils.server_utils import get_server_by_alias from src.utils.ssh import run_ssh_command @@ -27,7 +28,7 @@ async def overall_help(message: Message, state=None): @bot_action(r'meme admin docker ls', params=[ - ('alias', r'[\w\d]+') + BotParam('alias', r'[\w\d]+') ]) async def docker_ls(message: Message, alias, state=None): server = await get_server_by_alias(message, alias, ServerPermissions.DOCKER) @@ -60,7 +61,7 @@ async def docker_ls(message: Message, alias, state=None): @bot_action(r'meme admin docker clean', params=[ - ('alias', r'[\w\d]+') + BotParam('alias', r'[\w\d]+') ]) async def docker_clean(message: Message, alias, state=None): server = await get_server_by_alias(message, alias, ServerPermissions.DOCKER) @@ -78,8 +79,8 @@ async def docker_clean(message: Message, alias, state=None): @bot_action(r'meme admin docker rm', params=[ - ('alias', r'[\w\d]+'), - ('containers', r'.*') + BotParam('alias', r'[\w\d]+'), + BotParam('containers', r'.*') ]) async def docker_rm(message: Message, alias, containers, state=None): server = await get_server_by_alias(message, alias, ServerPermissions.DOCKER) @@ -97,8 +98,8 @@ async def docker_rm(message: Message, alias, containers, state=None): @bot_action(r'meme admin docker stop', params=[ - ('alias', r'[\w\d]+'), - ('containers', r'.*') + BotParam('alias', r'[\w\d]+'), + BotParam('containers', r'.*') ]) async def docker_rm(message: Message, alias, containers, state=None): server = await get_server_by_alias(message, alias, ServerPermissions.DOCKER) @@ -116,8 +117,8 @@ async def docker_rm(message: Message, alias, containers, state=None): @bot_action(r'meme admin docker up', params=[ - ('alias', r'[\w\d]+'), - ('containers', r'.*') + BotParam('alias', r'[\w\d]+'), + BotParam('containers', r'.*') ]) async def docker_rm(message: Message, alias, containers, state=None): server = await get_server_by_alias(message, alias, ServerPermissions.DOCKER) diff --git a/src/actions/interactive_session.py b/src/actions/interactive_session.py index 1110a4482d34130c7d799ae8916442b792a29825..bb1848952b6773b95465c1ddd0775ca12eb77621 100644 --- a/src/actions/interactive_session.py +++ b/src/actions/interactive_session.py @@ -22,7 +22,7 @@ from src.utils.decorators import bot_action from src.utils.server_utils import get_server_by_alias from src.utils.ssh import session_manager -logger = logging.getLogger() +logger = logging.getLogger(__name__) class UserStates(StatesGroup): @@ -75,7 +75,7 @@ async def choose_server(message: Message, state: FSMContext): async with state.proxy() as current_state: current_state['server_alias'] = chosen_server.server_alias current_state['session_id'] = await session_manager.add_connection(chosen_server) - logger.warning(current_state) + logger.warning(current_state['session_id']) await UserStates.interactive.set() await message.reply( f'Interactive session created for {server_alias}.\n' diff --git a/src/actions/permissions.py b/src/actions/permissions.py index 5405adb9fcb0b21284446fbbf42858523a2d0e49..2a712822d03bf014c027454d5507d799c613f5b8 100644 --- a/src/actions/permissions.py +++ b/src/actions/permissions.py @@ -6,10 +6,13 @@ from src.models.crud.server_crud import fn_update_server from src.models.server import ServerPermissions from src.settings import settings from src.utils.debug_mode import debug_message -from src.utils.decorators import bot_action +from src.utils.decorators import ( + bot_action, + BotParam +) from src.utils.server_utils import get_server_by_alias -logger = logging.getLogger() +logger = logging.getLogger(__name__) @bot_action(r'meme admin permissions( help)?$') @@ -32,8 +35,8 @@ async def list_permissions(message: Message, state=None): @bot_action(trigger_str=r'meme admin permissions grant', params=[ - ('server_alias', r'[\w\d]+'), - ('permission', r'\w+') + BotParam('server_alias', r'[\w\d]+'), + BotParam('permission', r'\w+') ]) async def change_permissions(message: Message, server_alias, permission, state=None): server = await get_server_by_alias(message, server_alias) diff --git a/src/actions/scp.py b/src/actions/scp.py new file mode 100644 index 0000000000000000000000000000000000000000..cd0f7f9b448c9cda74f4b492f4d72c908b65b4cf --- /dev/null +++ b/src/actions/scp.py @@ -0,0 +1,93 @@ +import logging + +from aiogram.dispatcher import FSMContext +from aiogram.dispatcher.filters.state import State, StatesGroup, default_state +from aiogram.types import Message, ParseMode, InputFile, ChatActions, ContentType + +from src.models.server import ServerPermissions +from src.utils.debug_mode import debug_message +from src.utils.decorators import ( + bot_action, + BotParam +) +from src.utils.server_utils import get_server_by_alias +from src.utils.ssh import download_file, upload_file + +logger = logging.getLogger(__name__) + + +class UploadingStates(StatesGroup): + start_uploading = State() + + +@bot_action(r'meme admin scp( help)?$') +async def scp_help(message: Message, state=None): + await message.reply( + '```\n' + '* upload {alias} {path} {user}?- Upload file on server as a user. ' + 'If user is null then the file will be uploaded as by root.\n\n' + '* download {alias} {file_path} {user}? - Download specific file from server as user.\n' + '```', + parse_mode=ParseMode.MARKDOWN + ) + + +@bot_action(r'meme admin scp download', + params=[ + BotParam('alias', r'[\w\d]+'), + BotParam('file_path', r'\.?[\/\w\-\d\.]+'), + BotParam('user', r'(\s+)?\.?[\/\w\-\d\.]+', optional=True) + ]) +async def download(message: Message, alias, file_path, user: str, state=None): + server = await get_server_by_alias(message, alias, ServerPermissions.INFO) + await ChatActions.upload_document() + user = user.strip() if user else None + filename = file_path.split("/")[-1] + async with download_file(server, file_path, user) as downloaded_file: + await message.reply_document(InputFile(downloaded_file, filename=filename), file_path) + + +@bot_action(r'meme admin scp upload', + params=[ + BotParam('alias', r'[\w\d]+'), + BotParam('file_path', r'\.?[\/\w\-\d\.]+'), + BotParam('user', r'(\s+)?\.?[\/\w\-\d\.]+', optional=True) + ]) +async def upload_start(message: Message, alias: str, file_path: str, user: str, state: FSMContext): + await debug_message( + message, + '```\n' + f'server: {alias}\n' + f'file: {file_path}\n' + f'user: {user.strip() if user else None}\n' + '```' + ) + await get_server_by_alias(message, alias, ServerPermissions.SCP) + await UploadingStates.start_uploading.set() + async with state.proxy() as current_state: + current_state['alias'] = alias + current_state['filename'] = file_path + current_state['user'] = user.strip() if user else None + await message.reply(f"Send a file in this chat to upload file on {alias} server") + + +@bot_action(filter_state=UploadingStates.start_uploading, content_types=ContentType.DOCUMENT) +async def upload(message: Message, state: FSMContext): + async with state.proxy() as current_state: + alias = current_state["user"] + user = current_state["user"] + filename = current_state["filename"] + server = await get_server_by_alias(message, current_state['alias'], ServerPermissions.SCP) + await default_state.set() + if not message.document: + await message.reply("This is not a document. Please attach a file.") + return + await debug_message( + message, + f'server: {alias}\n' + f'filename: {filename}\n' + f'user: {user}\n' + ) + async with upload_file(server, message.document, filename, user) as f: + await debug_message(message, f"File loaded at {f}") + await message.reply("File successfully uploaded") diff --git a/src/actions/server_management.py b/src/actions/server_management.py index aaf38eee91956c1ee6053787437256c9b785a730..084505ebd32201c449c315c73cc1ae38b4358617 100644 --- a/src/actions/server_management.py +++ b/src/actions/server_management.py @@ -11,10 +11,13 @@ from src.models.crud.server_crud import ( from src.models.server import ServerPermissions, Server from src.settings import settings from src.utils.debug_mode import debug_message -from src.utils.decorators import bot_action +from src.utils.decorators import ( + bot_action, + BotParam +) from src.utils.server_utils import get_server_by_alias -logger = logging.getLogger() +logger = logging.getLogger(__name__) @bot_action(f'meme admin servers( help)?$') @@ -32,8 +35,8 @@ async def servers_help(message: Message, state=None): @bot_action(trigger_str=r'meme admin servers add', params=[ - ('address', r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}(\:(?P<port>\d+))?'), - ('alias', r'[\w\d]+') + BotParam('address', r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}(\:(?P<port>\d+))?'), + BotParam('alias', r'[\w\d]+') ]) async def add_server(message: Message, address, alias, port, state=None): await debug_message(message, @@ -61,8 +64,8 @@ async def add_server(message: Message, address, alias, port, state=None): @bot_action(trigger_str=r'meme admin servers rename', params=[ - ('old_alias', r'[\w\d]+'), - ('new_alias', r'[\w\d]+') + BotParam('old_alias', r'[\w\d]+'), + BotParam('new_alias', r'[\w\d]+') ]) async def rename_server(message: Message, old_alias, new_alias, state=None): server = await get_server_by_alias(message, old_alias) @@ -84,7 +87,7 @@ async def rename_server(message: Message, old_alias, new_alias, state=None): @bot_action(trigger_str=r'meme admin servers delete', params=[ - ('server_alias', r'[\w\d]+') + BotParam('server_alias', r'[\w\d]+') ]) async def delete_server(message: Message, server_alias, state=None): server = await get_server_by_alias(message, server_alias) diff --git a/src/actions/stats.py b/src/actions/stats.py index 0f9226ccf62db218fd95b65f4543eb43002b0bb4..94a38b650b7502d78be2a8990ff1f5d6e9bec22e 100644 --- a/src/actions/stats.py +++ b/src/actions/stats.py @@ -3,11 +3,14 @@ import logging from aiogram.types import Message, ParseMode from src.utils.debug_mode import debug_message -from src.utils.decorators import bot_action +from src.utils.decorators import ( + bot_action, + BotParam +) from src.utils.server_utils import get_server_by_alias from src.utils.ssh import run_ssh_command -logger = logging.getLogger() +logger = logging.getLogger(__name__) @bot_action(r'meme admin stats( help)?$') @@ -24,7 +27,7 @@ async def overall_help(message: Message, state=None): @bot_action(r'meme admin stats cpu', params=[ - ('alias', r'[\w\d]+') + BotParam('alias', r'[\w\d]+') ]) async def cpu_usage(message: Message, alias: str, state=None): server = await get_server_by_alias(message, alias) @@ -62,7 +65,7 @@ async def cpu_usage(message: Message, alias: str, state=None): @bot_action(r'meme admin stats ram', params=[ - ('alias', r'[\w\d]+') + BotParam('alias', r'[\w\d]+') ]) async def ram_usage(message: Message, alias: str, state=None): server = await get_server_by_alias(message, alias) @@ -94,7 +97,7 @@ async def ram_usage(message: Message, alias: str, state=None): @bot_action(r'meme admin stats mem', params=[ - ('alias', r'[\w\d]+') + BotParam('alias', r'[\w\d]+') ]) async def mem(message: Message, alias: str, state=None): server = await get_server_by_alias(message, alias) diff --git a/src/actions/systemctl.py b/src/actions/systemctl.py index 8a922405136dfe1dedc2b84aa89c7bb16c4c9518..14dc904de7425595efe6b4c4c2759e7f95cde909 100644 --- a/src/actions/systemctl.py +++ b/src/actions/systemctl.py @@ -6,7 +6,8 @@ from aiogram.types import Message, ParseMode from src.models.server import ServerPermissions from src.utils.debug_mode import debug_message from src.utils.decorators import ( - bot_action + bot_action, + BotParam ) from src.utils.server_utils import get_server_by_alias from src.utils.ssh import run_ssh_command @@ -30,8 +31,8 @@ async def overall_help(message: Message, state=None): @bot_action(r'meme admin sys stats', params=[ - ('alias', r'[\w\d]+'), - ('service', r'[\w\d]+\.service') + BotParam('alias', r'[\w\d]+'), + BotParam('service', r'[\w\d]+\.service') ]) async def status(message: Message, alias, service, state): server = await get_server_by_alias(message, @@ -57,8 +58,8 @@ async def status(message: Message, alias, service, state): @bot_action(r'meme admin sys start', params=[ - ('alias', r'[\w\d]+'), - ('service', r'[\w\d]+\.service') + BotParam('alias', r'[\w\d]+'), + BotParam('service', r'[\w\d]+\.service') ]) async def start(message: Message, alias, service, state): server = await get_server_by_alias(message, @@ -70,8 +71,8 @@ async def start(message: Message, alias, service, state): @bot_action(r'meme admin sys stop', params=[ - ('alias', r'[\w\d]+'), - ('service', r'[\w\d]+\.service') + BotParam('alias', r'[\w\d]+'), + BotParam('service', r'[\w\d]+\.service') ]) async def stop(message: Message, alias, service, state): server = await get_server_by_alias(message, @@ -83,8 +84,8 @@ async def stop(message: Message, alias, service, state): @bot_action(r'meme admin sys restart', params=[ - ('alias', r'[\w\d]+'), - ('service', r'[\w\d]+\.service') + BotParam('alias', r'[\w\d]+'), + BotParam('service', r'[\w\d]+\.service') ]) async def restart(message: Message, alias, service, state): server = await get_server_by_alias(message, @@ -96,8 +97,8 @@ async def restart(message: Message, alias, service, state): @bot_action(r'meme admin sys cat', params=[ - ('alias', r'[\w\d]+'), - ('service', r'[\w\d]+\.service') + BotParam('alias', r'[\w\d]+'), + BotParam('service', r'[\w\d]+\.service') ]) async def cat(message: Message, alias, service, state): server = await get_server_by_alias(message, diff --git a/src/models/server.py b/src/models/server.py index b00699b76b352d662b48afcb358cb2b21a541479..370afec4f84857bb5de2cb72ba556426c6b92684 100644 --- a/src/models/server.py +++ b/src/models/server.py @@ -7,6 +7,7 @@ from src.utils.enum_utils import AutoEnum class ServerPermissions(AutoEnum): INFO = "Users can read server info" + SCP = "Users can download and upload files to server" DOCKER = "Users can manage docker daemon" SERVICES = "Users can manage systemctl services" RCE = "Users can execute remote commands on this server" diff --git a/src/utils/debug_mode.py b/src/utils/debug_mode.py index 6422a73bb01074d6f6b67f75a23f7bc01785d7d9..8855a36adfbb57f4cea0a6b0debaa4429d0a3316 100644 --- a/src/utils/debug_mode.py +++ b/src/utils/debug_mode.py @@ -4,7 +4,7 @@ from aiogram.types import Message, ParseMode from src.settings import settings -logger = logging.getLogger() +logger = logging.getLogger(__name__) async def debug_message(message: Message, text: str): diff --git a/src/utils/decorators.py b/src/utils/decorators.py index c106b1129aefbd0c2191aade62f2280f747f270e..40849e511c78ebbe674057b7e6f31f606f736ba1 100644 --- a/src/utils/decorators.py +++ b/src/utils/decorators.py @@ -8,7 +8,19 @@ from aiogram.types import Message, ParseMode, ReplyKeyboardRemove from src.settings import settings -logger = logging.getLogger() +logger = logging.getLogger(__name__) + + +class BotParam(object): + def __init__(self, name, regex, optional=False): + self.name = name + self.regex = regex + self.optional = optional + + def get_regex(self): + option = "?" if self.optional else "" + sep = "" if self.optional else " " + return f"{sep}(?P<{self.name}>{self.regex}){option}" def cool_response_exception(f: Callable): @@ -34,7 +46,8 @@ def cool_response_exception(f: Callable): return decorated_func -def bot_action(trigger_str: str = None, params: List[Tuple] = None, filter_state: State = None) -> Callable: +def bot_action(trigger_str: str = None, params: List[Tuple] = None, filter_state: State = None, + content_types=None) -> Callable: """ Mark function as bot message handler. params -> is just an array of tuples. @@ -44,15 +57,18 @@ def bot_action(trigger_str: str = None, params: List[Tuple] = None, filter_state :param trigger_str: message that triggers bot :param params: parameters in message :param filter_state: trigger state if you use Finite State Machine mechanism + :param content_types: Which content type bot is waiting for :return: decorator """ def decor(f: Callable): - regexp_params = ' '.join([f'(?P<{name}>{pattern})' for name, pattern in params]) if params else '' - filter_regexp = f'^{trigger_str} {regexp_params}'.strip() + regexp_params = ''.join(map(lambda x: x.get_regex(), params)) if params else '' + filter_regexp = f'^{trigger_str}{regexp_params}'.strip() @wraps(f) - @settings.dispatcher.message_handler(regexp=filter_regexp if trigger_str else None, state=filter_state) + @settings.dispatcher.message_handler(regexp=filter_regexp if trigger_str else None, + state=filter_state, + content_types=content_types) @cool_response_exception async def modified_function(message: Message, state=None): additional_values = {} diff --git a/src/utils/ssh.py b/src/utils/ssh.py index 7452cbc280354754f53c4104a3fe827a05e989fd..b02f219f544b26d7074d28f03d5db4ee62a90663 100644 --- a/src/utils/ssh.py +++ b/src/utils/ssh.py @@ -1,27 +1,42 @@ import asyncio import logging +import os import uuid +from contextlib import asynccontextmanager import asyncssh +from aiogram.types import Document from src.models import Server -logger = logging.getLogger() +logger = logging.getLogger(__name__) async def open_ssh_session(server: Server) -> asyncssh.SSHClientProcess: + """ + Open bash process over ssh. + + :param server: server where pipe needs to be opened + :return: opened process + """ connection = await asyncssh.connect(server.server_address, server.server_port, known_hosts=None) process = await connection.create_process('/bin/bash') return process async def collect_output(out_pipe, timeout): + """ + Collect output from command. + output will be collected from out_pipe. + + :param out_pipe: async io channel + :param timeout: timeout to read from channel + :return: + """ res = [] try: line = await asyncio.wait_for(out_pipe.readline(), timeout) - res.append(line) while line: - logger.info(line) res.append(await asyncio.wait_for(out_pipe.readline(), timeout)) except asyncio.exceptions.TimeoutError: return ''.join(res).strip() @@ -31,34 +46,110 @@ async def collect_output(out_pipe, timeout): async def run_interactive_command(command: str, process: asyncssh.SSHClientProcess, timeout=0.5) -> str: + """ + Run command on running ssh session. + Command will be executed on server with provided process. + Timeout used to collect output from commands. + + :param command: command yo execute + :param process: where ot execute + :param timeout: receiving timeout + :return:stdout or stderr as exception + """ process.stdin.write(command + '\n') stdout = await collect_output(process.stdout, timeout) stderr = await collect_output(process.stderr, timeout) - logger.debug(f"STDOUT: {stdout}") - logger.debug(f"STDERR: {stderr}") if stderr: raise Exception(stderr) return stdout async def run_ssh_command(server: Server, command: str) -> str: + """ + Run command on remote server. + Passed command will be executed as root. + + :param server: remote server + :param command: command to execute + :return: output. + """ process = await open_ssh_session(server) res = await run_interactive_command(command, process) process.close() return res +@asynccontextmanager +async def download_file(server: Server, file_path: str, username: str) -> str: + """ + Download file from server and store it as random file in tmp. + When context is over, than it will remove file from tmp. + + :param server: remote server. + :param file_path: file path on remote server + :param username: as who you want to log in + :return: path to downloaded file + """ + load_path = f'/tmp/{uuid.uuid4()}' + async with asyncssh.connect(server.server_address, + server.server_port, + known_hosts=None, + username=username) as conn: + await asyncssh.scp((conn, file_path), load_path) + yield load_path + os.remove(load_path) + + +@asynccontextmanager +async def upload_file(server: Server, document: Document, file_path: str, username: str): + """ + Upload message document to a server. + + :param server: server where to upload + :param document: amessage attached document + :param file_path: remote file path to store file + :param username: user to login as + :return: + """ + load_path = f'/tmp/{uuid.uuid4()}' + await document.download(load_path) + async with asyncssh.connect(server.server_address, + server.server_port, + known_hosts=None, + username=username) as conn: + await asyncssh.scp(load_path, (conn, file_path)) + yield load_path + os.remove(load_path) + + class SessionManager(object): + """ + Class to manage all connection to other servers over ssh. + """ def __init__(self): self.__connections = dict() async def add_connection(self, server: Server) -> str: + """ + Add connection some server and assign uuid to it. + This method will open chanel and return uuid of session. + + :param server: server to open connection at + :return: valid uuid4 string + """ rand_uuid = str(uuid.uuid4()) session = await open_ssh_session(server) self.__connections[rand_uuid] = session return rand_uuid async def run_command(self, connection_id: str, command: str): + """ + Run command in opened session by id. + + :param connection_id: valid uuid of opened session + :param command: command to run over ssh + :return: command output + """ return await run_interactive_command(command, self.__connections[connection_id]) def close(self, connection_id: str):