mosh3: a mosh helper that reuses ssh connections (ControlMaster)

mosh3.py

#!/usr/bin/env python3

# inspired by
# https://github.com/mobile-shell/mosh/issues/24#issuecomment-201893250

import sys
import os
import subprocess

def main():
  args = sys.argv[1:]
  try:
    pos = args.index('--')
    ssh_args = args[:pos]
    mosh_args = args[pos+1:]
  except ValueError:
    ssh_args = args
    mosh_args = []

  cmd = ['ssh', '-T'] + ssh_args
  code = '''\
{
  . /etc/locale.conf >&2 && export LANG
} || export LANG=zh_CN.UTF-8
env -u SHLVL mosh-server -- %s || exit
echo -n 'IP '
# curl -sS https://ipinfo.io/ip
curl -sS4 http://ip.sb/
''' % subprocess.list2cmdline(mosh_args)

  p = subprocess.run(
    cmd,
    stdout = subprocess.PIPE,
    # stderr = subprocess.DEVNULL,
    input = code,
    universal_newlines = True,
    check = True,
  )
  port = key = ip = None
  for line in p.stdout.splitlines():
    if line.startswith('MOSH CONNECT '):
      port, key = line.split()[2:]
    elif line.startswith('IP '):
      ip = line.split()[1]
      break

  if not port or not key:
    sys.exit('Did not find mosh-server info')
  elif not ip:
    sys.exit('Did not find server IP')

  env = os.environ.copy()
  env['MOSH_KEY'] = key
  os.execvpe('mosh-client', ['mosh-client', ip, port], env)

if __name__ == '__main__':
  main()