Source

beanstalk-stack / beanstalk / bs_ssh_entry

Full commit
#!/usr/bin/env python

import os
import re
#import shlex
import sys

valid_commands = [r'^rsync --server .*', r'^bsjack server\..+', '^bsjack info\..+']
fabric_pattern = re.compile(r'^/bin/bash -l -c "(?P<command>.*)"$')


def fabric_command(command):
    match = fabric_pattern.match(command)
    return match.groupdict()['command'].replace('\\', '') if match else command


def valid_command(command):
    for valid_command in valid_commands:
        if re.compile(valid_command).match(command) is not None:
            # rsync path check
            #if command.startswith('rsync --server'):
            #    target_path = shlex.split(command)[-1]
            #    if not target_path.startswith('/var/beanstalk'):  # TODO: Set this path by settings
            #        return False
            return True
    return False


def main():
    command = os.environ.get('SSH_ORIGINAL_COMMAND', None)
    if command is None:
        print 'You cannot use SSH as beanstalk directly.'
        sys.exit(1)

    command = fabric_command(command)
    if not valid_command(command):
        print 'Invalid command: {command}'.format(command=command)
        sys.exit(2)

    os.system(command)


if __name__ == '__main__':
    main()