Show More
ssh_wrapper_v1.py
72 lines
| 2.9 KiB
| text/x-python
|
PythonLexer
r5608 | # Copyright (C) 2016-2024 RhodeCode GmbH | |||
r5314 | # | |||
# This program is free software: you can redistribute it and/or modify | ||||
# it under the terms of the GNU Affero General Public License, version 3 | ||||
# (only), as published by the Free Software Foundation. | ||||
# | ||||
# 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 Affero General Public License | ||||
# along with this program. If not, see <http://www.gnu.org/licenses/>. | ||||
# | ||||
# This program is dual-licensed. If you wish to learn more about the | ||||
# RhodeCode Enterprise Edition, including its added features, Support services, | ||||
# and proprietary license terms, please see https://rhodecode.com/licenses/ | ||||
import os | ||||
import sys | ||||
import time | ||||
import logging | ||||
import click | ||||
from rhodecode.lib.pyramid_utils import bootstrap | ||||
from rhodecode.lib.statsd_client import StatsdClient | ||||
from .backends import SshWrapper | ||||
r5316 | from .utils import setup_custom_logging | |||
r5314 | ||||
log = logging.getLogger(__name__) | ||||
@click.command() | ||||
@click.argument('ini_path', type=click.Path(exists=True)) | ||||
@click.option( | ||||
'--mode', '-m', required=False, default='auto', | ||||
type=click.Choice(['auto', 'vcs', 'git', 'hg', 'svn', 'test']), | ||||
help='mode of operation') | ||||
@click.option('--user', help='Username for which the command will be executed') | ||||
@click.option('--user-id', help='User ID for which the command will be executed') | ||||
@click.option('--key-id', help='ID of the key from the database') | ||||
@click.option('--shell', '-s', is_flag=True, help='Allow Shell') | ||||
@click.option('--debug', is_flag=True, help='Enabled detailed output logging') | ||||
def main(ini_path, mode, user, user_id, key_id, shell, debug): | ||||
setup_custom_logging(ini_path, debug) | ||||
command = os.environ.get('SSH_ORIGINAL_COMMAND', '') | ||||
if not command and mode not in ['test']: | ||||
raise ValueError( | ||||
'Unable to fetch SSH_ORIGINAL_COMMAND from environment.' | ||||
'Please make sure this is set and available during execution ' | ||||
'of this script.') | ||||
connection_info = os.environ.get('SSH_CONNECTION', '') | ||||
time_start = time.time() | ||||
with bootstrap(ini_path, env={'RC_CMD_SSH_WRAPPER': '1'}) as env: | ||||
r5325 | settings = env['registry'].settings | |||
r5314 | statsd = StatsdClient.statsd | |||
try: | ||||
ssh_wrapper = SshWrapper( | ||||
command, connection_info, mode, | ||||
r5325 | user, user_id, key_id, shell, ini_path, settings, env) | |||
r5314 | except Exception: | |||
log.exception('Failed to execute SshWrapper') | ||||
sys.exit(-5) | ||||
return_code = ssh_wrapper.wrap() | ||||
operation_took = time.time() - time_start | ||||
if statsd: | ||||
operation_took_ms = round(1000.0 * operation_took) | ||||
statsd.timing("rhodecode_ssh_wrapper_timing.histogram", operation_took_ms, | ||||
use_decimals=False) | ||||
sys.exit(return_code) | ||||