ssh-wrapper

#!/usr/bin/python3
# vim: fileencoding=utf-8
#
# The Qubes OS Project, https://www.qubes-os.org/
#
# Copyright (C) 2016 Marek Marczykowski-Górecki
#                                       <marmarek@invisiblethingslab.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# 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 General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
import argparse
import subprocess
import sys


def ssh(args):
    # WARNING: this doesn't support interactive communication, just sending
    # script through stdin in one go and return its output
    start_string = "/bin/sh  << 'EOF'\n"
    end_string = '\nEOF'
    if args[1] == '/bin/sh':
        stdin_data = sys.stdin.buffer.read()
    elif args[1].startswith(start_string) and args[1].endswith(end_string):
        stdin_data = args[1].encode('ascii', 'strict')
        if stdin_data.find((end_string + '\n').encode()) != -1:
            raise ValueError('input is not single heredoc')
        stdin_data = stdin_data[len(start_string):1 - len(end_string)]
    else:
        raise ValueError('do not know how to handle our input')

    # Currently when managing VM through  qubesctl  (so, wrapped  salt-ssh ),
    # it checks for  scp binary presence in the target VM. In our case it
    # doesn't make sense, since our wrapper uses qubes.Filecopy qrexec
    # service instead - so scp isn't required at all.

    # This check is hardcoded in salt thin minion wrapper, so probably not
    # easy to disable. But it is possible to simulate scp binary
    # presence, by adding additional directory to PATH with a dummy scp file.
    # So lets go this way, instead of requiring scp being installed in all
    # the templates.
    stdin_data = (
        b"mkdir -p /run/salt-shim-sandbox\n"
        b"ln -sf /bin/true /run/salt-shim-sandbox/scp\n"
        b"export PATH=\"$PATH:/run/salt-shim-sandbox\"\n"
        + stdin_data
    )
    p = subprocess.Popen(['qrexec-client-vm', args[0], 'qubes.VMRootShell'],
        stdin=subprocess.PIPE)
    p.communicate(stdin_data)
    # if qubes.VMRootShell service not supported, fallback to qubes.VMShell and
    # hope it will have appropriate permissions
    if p.returncode == 127:
        p = subprocess.Popen(['qrexec-client-vm', args[0], 'qubes.VMShell'],
            stdin=subprocess.PIPE)
        p.communicate(stdin_data)
    return p.returncode

def scp(args):
    assert len(args) == 2 and args[1].count(':') == 1
    src_path = args[0]
    (dst_host, dst_path) = args[1].split(':')
    source_content = open(src_path, 'rb').read()
    p = subprocess.Popen(['qrexec-client-vm', dst_host, 'qubes.VMRootShell'],
        stdin=subprocess.PIPE)
    p.communicate('cat > "{}"\n'.format(dst_path).encode() + source_content)
    # if qubes.VMRootShell service not supported, fallback to qubes.VMShell and
    # hope it will have appropriate permissions
    if p.returncode == 127:
        p = subprocess.Popen(['qrexec-client-vm', dst_host, 'qubes.VMShell'],
            stdin=subprocess.PIPE)
        p.communicate('cat > "{}"\n'.format(dst_path).encode() + source_content)
    if p.returncode != 0:
        raise RuntimeError('Failed to write target file {}'.format(dst_path))
    return 0


def parse_opts(args):
    parser = argparse.ArgumentParser()
    parser.add_argument('-o', action='append', nargs=1)
    parser.add_argument('--version', '-V', action='store_true')
    (opts, args) = parser.parse_known_args(args)
    if opts.version:
        print('OpenSSH_6.6.1p1 qubes.VMShell wrapper, '
              'OpenSSL 1.0.1k-fips 8 Jan 2015')
        sys.exit(0)
    return args


def main(args=None):
    args = parse_opts(args)
    if sys.argv[0].endswith('ssh'):
        return ssh(args)
    elif sys.argv[0].endswith('scp'):
        return scp(args)
    else:
        raise RuntimeError('Unsupported program {} called'.format(sys.argv[0]))

if __name__ == '__main__':
    sys.exit(main())