#!/usr/bin/env python
# @package      hubzero-mw2-exec-virtualssh
# @file         /usr/bin/vssh_exec_proxy
# @copyright    Copyright (c) 2016-2020 The Regents of the University of California.
# @license      http://opensource.org/licenses/MIT MIT
#
# Based on prior work by Richard L. Kennell
#
# Copyright (c) 2016-2020 The Regents of the University of California.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# HUBzero is a registered trademark of The Regents of the University of California.
#

"""This runs on the execution host.  Its goal is to connect to an SSH server
running inside a tool container.   
1. Is there an SSH server?  If not call the script /usr/bin/vsshd_start
It will do the following:
Start an SSH server inside the container, listening on port 2200, running as the user.  
Use the configuration file /etc/mw-virtualssh/sshd_config.
Create SSH keys.  Copy public key to "/ssh/authorized_keys"
Call ensure-known-host to add container IP and SSH key to /etc/ssh/ssh_known_hosts

2. SSH into the container, redirecting stdin and out


To debug virtual SSH, go inside the container and start the server manually with:
su <username>
/usr/sbin/sshd -f /ssh/sshd_config -d -d

and from the execution host, run:
ssh -F /etc/mw-virtualssh/exec_ssh_config -i /etc/mw-virtualssh/ssh_guest_key -t -p 2200 username@ip.ad.dre.ss -v -v -v 
where ip.ad.dre.ss is the container's private IP and username is the user account for whom the tool session is running

"""
import os
import sys
import socket
import subprocess
import re
import time
import syslog

from hubzero.mw.constants import EXEC_CONFIG_FILE, VIRTUALSSH_K, USER_REGEXP
from hubzero.mw.container import Container, make_Container

CONTAINER_CONF = {}
execfile(EXEC_CONFIG_FILE)

VSSH_PROXY_LOG = '/var/log/vssh-exec-proxy/vssh-exec-proxy.log'
idfile = "/etc/mw-virtualssh/ssh_guest_key"
ENSURE_PATH = "/usr/bin/ensure-known-host"
SSH_CONFIG_PATH = '/etc/mw-virtualssh/exec_ssh_config'
# vsshd_start creates the ssh keys if necessary, and copies them to the container
# it also starts the ssh server inside the container with the user's id
VSSHD_START_PATH = '/usr/bin/vsshd_start'

DEBUG = False
DOCKER_EXEC = True
if DEBUG:
  print "using a container of class: " + CONTAINER_CONF['class']
  print os.environ
  print sys.argv
  print CONTAINER_CONF['class']
  if not os.environ.has_key('SSH_TTY'):
    print "option -t wasn't used"

try:
  dispnum = int(sys.argv[1])
except ValueError:
  if DEBUG:
    print "Unable to determine display number"
  syslog.syslog("vssh_exec_proxy: Unable to determine display number")
  sys.exit(1)

try:
  username = sys.argv[2]
except (TypeError, ValueError):
  if DEBUG:
    print "Unable to determine username"
  syslog.syslog("vssh_exec_proxy: Unable to determine username")
  sys.exit(1)

m = re.match(USER_REGEXP, username)
if m is None:
  if DEBUG:
    print "Bad username"
  syslog.syslog("vssh_exec_proxy: Bad username '%s'" % username)
  sys.exit(1)

if dispnum <= 0 or dispnum > 65535:
  if DEBUG:
    print "Bad container id"
  syslog.syslog("vssh_exec_proxy: Bad container id '%d'" % dispnum)
  sys.exit(1)

if DEBUG:
  print "input validation successful"

command = str.join(' ',sys.argv[3:])
if DEBUG:
  print "requested command to run in container is '%s'" % command

# access control should have been done by front virtualssh proxy, we have no way of doing access control
syslog.syslog("vssh_exec_proxy: connect to container %d with username '%s'" % (dispnum, username))
  
# use Container class to make sure IP address is always calculated the same way
# (disp, machine_number, overrides={})
# vps = Container(dispnum, 0, CONTAINER_CONF)
# use make_Container so that the IP address is calculated according to the correct deployment model and virtualization tech
# but vps.veaddr gives IP address of services Docker container instead of tool container!
vps = make_Container(dispnum, 0, CONTAINER_CONF)
if CONTAINER_CONF['class'] == 'ContainerDocker' or CONTAINER_CONF['class'] == 'ContainerAWS':
  addr = vps.tool_container_IP
  
  # Check if container exists, wait several seconds in case it is being started
  i = 0
  containerID = ""
  check_args = ['/usr/bin/docker', 'ps', '-q', '-f', 'name=^%d.tool' % dispnum]
  while i <100 and containerID == "":
    process = subprocess.Popen(check_args, stdout = subprocess.PIPE)
    (containerID, dummy) = process.communicate()
    if containerID == "":
      time.sleep(1)
      i+=1
  if i >= 100:
    print "Tool container not found"
    sys.exit(1)
  if DEBUG:
    print "containerID is " + containerID

  # if allowed, and not doing X11 forwarding, run docker exec /bin/bash, as the user because it's faster
  if DOCKER_EXEC and not os.environ.has_key('DISPLAY'):
    if DEBUG:
      print "docker environment and no DISPLAY specified"
    args = []
    # The following environment variables are available only if sshd is running with the options
    # AcceptEnv SESSION SESSIONDIR RESULTSDIR
    if os.environ.has_key('SESSION'):
      args += ['-e', 'SESSION=%s' % os.environ['SESSION']]
    if os.environ.has_key('SESSIONDIR'):
      args += ['-e', 'SESSIONDIR=%s' % os.environ['SESSIONDIR']]
    if os.environ.has_key('RESULTSDIR'):
      args += ['-e', 'RESULTSDIR=%s' % os.environ['RESULTSDIR']]

    args += ['-e', 'VDISPLAY=%s:0.0' % addr]
    args += ['--user', username]
    args += ['-e', "USER=%s" % username]
    args += ['-e', "PATH=%s" % "/bin:/usr/bin:/usr/bin/X11:/sbin:/usr/sbin"]    
    args += ['%d.tool' % dispnum]
    if len(command) > 0:
      if DEBUG:
        print "trying docker exec"
      args = ['/usr/bin/docker', 'exec'] + args + ['/bin/bash', '-c', command]
    else:
      if DEBUG:
        print "no command so trying docker exec -it if SSH_TTY is defined"
      # Check SSH_TTY to avoid Docker returning error "The input device is not a TTY"
      if os.environ.has_key('SSH_TTY'):
        args = ['/usr/bin/docker', 'exec', '-it'] + args + ['/bin/bash', '-l']
      else:
        if DEBUG:
          print "SSH_TTY not defined so entering via SSH server in container instead of using docker exec"
        args = []
    if len(args) > 0:
      os.execve("/usr/bin/docker", args, os.environ)
  elif DEBUG:
    print "Docker container but there is a DISPLAY environment variable %s, skipped docker exec and using instead an SSH server inside container" % os.environ['DISPLAY']
else:
  addr = vps.veaddr

# all arguments must be strings, so use "%d" % dispnum
# try several times to avoid "No passwd entry for user 'xyz'" messages in new containers
i=0
process = subprocess.Popen([VSSHD_START_PATH, "%d" % dispnum, username], stderr = subprocess.PIPE)
process.communicate()
rcode = process.returncode
while i <100 and rcode != 0:
  process = subprocess.Popen([VSSHD_START_PATH, "%d" % dispnum, username], stderr = subprocess.PIPE)
  process.communicate()
  rcode = process.returncode
  # print "sleeping 1 second"
  time.sleep(1)
  i+=1

if rcode != 0:
  syslog.syslog("vssh_exec_proxy: Unable to start or setup an SSH server inside the tool session (is the session still running?)")
  print "Unable to start or setup an SSH server inside the tool session (is the session still running?)"
  process = subprocess.Popen([VSSHD_START_PATH, "%d" % dispnum, username], stderr = subprocess.PIPE)
  syslog.syslog("vssh_exec_proxy: %s %s" % process.communicate())
  sys.exit(1)

# messing with the known_hosts doesn't help communications within the execution host
# so we have "StrictHostKeyChecking no"
# without that we'd always be copying the key with full trust anyway
# subprocess.check_call([ENSURE_PATH, addr, '2200'])

flags=[]
# if len(command) == 0 or os.environ.has_key('SSH_TTY'):
# two -t options force tty allocation but then ctrl-D doesn't work to disconnect
if os.environ.has_key('SSH_TTY'):
  flags += [ '-t' ]
if os.environ.has_key('SSH_AUTH_SOCK'):
  flags += [ '-A' ]
if os.environ.has_key('DISPLAY'):
  flags += [ '-X', '-Y' ]

if DEBUG:
  print "ssh-client command is '%s'" % command
  print "flags are ",
  print flags
  print "/usr/bin/ssh %s %s %s %s %s %s %s %s %s %s" % ('ssh', '-F', SSH_CONFIG_PATH, '-i', idfile,flags,'-p', '2200', '%s@%s' % (username,addr), command)

#os.system('env')
#time.sleep(.5)
pid = os.fork()
if pid == 0:
  os.environ['VDISPLAY'] = '%s:0.0' % addr
  os.execve("/usr/bin/ssh",
            ['ssh', '-F', SSH_CONFIG_PATH, '-i', idfile ] + flags +
            [ '-p', '2200', '%s@%s' % (username,addr), command],
            os.environ)
  if DEBUG:
    print "ssh command done"
else:
  os.wait()

sys.exit(0)

