Teach mbash to route SSH connections to the user's default pool

Author: quentinmit

server/common/oursrc/accountadm/mbash.in

@@ -1,3 +1,98 @@
-#!/bin/sh
+#!/usr/bin/env python
 
-exec @bash_path@ --rcfile /usr/local/etc/mbashrc "$@"
+from __future__ import (absolute_import, division, print_function)
+
+import os
+import sys
+import getpass
+import subprocess
+import ldap
+import ldap.filter
+
+def get_pool(username):
+    ldap_uri = ldap.get_option(ldap.OPT_URI)
+
+    ll = ldap.initialize(ldap_uri)
+
+    users = ll.search_s(
+        'dc=scripts,dc=mit,dc=edu',
+        ldap.SCOPE_SUBTREE,
+        ldap.filter.filter_format('(&(objectClass=posixAccount)(uid=%s))', [username]),
+        [],
+    )
+    if not users:
+        return None, None
+    user_dn = users[0][0]
+
+    pool_ips = set()
+    vhost_pools = {}
+    for dn, attrs in ll.search_s(
+        'dc=scripts,dc=mit,dc=edu',
+        ldap.SCOPE_SUBTREE,
+        ldap.filter.filter_format('(&(objectClass=scriptsVhost)(scriptsVhostAccount=%s))', [user_dn]),
+        ['scriptsVhostName', 'scriptsVhostPoolIPv4'],
+    ):
+        vhost_pools[attrs['scriptsVhostName'][0]] = attrs['scriptsVhostPoolIPv4'][0]
+        pool_ips.add(attrs['scriptsVhostPoolIPv4'][0])
+
+    pool_names = {}
+    for dn, attrs in ll.search_s(
+            'dc=scripts,dc=mit,dc=edu',
+            ldap.SCOPE_SUBTREE,
+            '(&(objectClass=scriptsVhostPool)(|'+''.join(ldap.filter.filter_format('(scriptsVhostPoolIPv4=%s)', [ip]) for ip in pool_ips)+'))',
+            ['cn', 'scriptsVhostPoolIPv4'],
+    ):
+        pool_names[attrs['scriptsVhostPoolIPv4'][0]] = attrs['cn'][0]
+
+    main_pool = vhost_pools.get(username + '.scripts.mit.edu')
+    other_pools = None
+    if len(pool_ips) > 1:
+        other_pools = sorted((pool_names.get(pool, pool), vhost) for vhost, pool in vhost_pools.items())
+    return main_pool, other_pools
+
+def should_forward():
+    ssh_connection = os.environ.get('SSH_CONNECTION')
+    if not ssh_connection:
+        return False
+    _, _, laddr, _ = ssh_connection.split(' ')
+    try:
+        with open('/etc/scripts/mbash-vips') as f:
+            if laddr in [l.strip() for l in f]:
+                return True
+    except IOError:
+        return False
+    return False
+
+def has_pool(ip):
+    return len(subprocess.check_output(['/sbin/ip', 'addr', 'show', 'to', ip])) > 0
+
+def maybe_forward():
+    if not should_forward():
+        return
+    command = None
+    if len(sys.argv) == 3 and sys.argv[1] == '-c':
+        command = sys.argv[2]
+    elif len(sys.argv) != 1:
+        print("Unexpected shell invocation; not forwarding.", file=sys.stderr)
+        return
+    user = getpass.getuser()
+    main_pool, other_pools = get_pool(user)
+    forward = main_pool and not has_pool(main_pool)
+    if forward:
+        # TODO: Check if we're already on the right server.
+        print("Forwarding to the server for %s.scripts.mit.edu." % (user,), file=sys.stderr)
+    if other_pools:
+        print("Your account has virtual hosts on multiple server pools; to connect to a server for a particular host, connect to a specific server:", file=sys.stderr)
+        print(file=sys.stderr)
+        for name, vhost in other_pools:
+            print("%s - ssh %s" % (vhost, name), file=sys.stderr)
+        print(file=sys.stderr)
+    if forward:
+        args = ['ssh', main_pool, '--']
+        if command is not None:
+            args.append(command)
+        os.execv('/usr/bin/ssh', args)
+
+maybe_forward()
+
+os.execv("@bash_path@", ["bash", "--rcfile", "/usr/local/etc/mbashrc"] + sys.argv[1:])