fishdescriptor/fishdescriptor

   1 #!/usr/bin/python3
   2
   3 # This file is part of chiark-utils, a collection of useful programs
   4 # used on chiark.greenend.org.uk.
   5 #
   6 # This file is:
   7 #  Copyright 2018 Citrix Systems Ltd
   8 #
   9 # This is free software; you can redistribute it and/or modify it under the
  10 # terms of the GNU General Public License as published by the Free Software
  11 # Foundation; either version 3, or (at your option) any later version.
  12 #
  13 # This is distributed in the hope that it will be useful, but WITHOUT ANY
  14 # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  15 # FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
  16 # details.
  17 #
  18 # You should have received a copy of the GNU General Public License along
  19 # with this program; if not, consult the Free Software Foundation's
  20 # website at www.fsf.org, or the GNU Project website at www.gnu.org.
  21
  22 import sys
  23 import fishdescriptor.fish
  24 import optparse
  25 import re
  26 import subprocess
  27 import socket
  28 import os
  29
  30 donor = None
  31
  32 usage = '''fishdescriptor [-p|--pid] <pid> <action>... [-p|--pid <pid> <action>...]
  33
  34 <action>s
  35   [<here-fd>=]<there-fd>
  36           fish the openfile referenced by descriptor <there-fd> in
  37           (the most recent) <pid> and keep a descriptor onto it;
  38           and, optionally, give it the number <here-fd> for exec
  39   exec <program> [<arg>...]
  40           execute a process with each specified <here>
  41           as an actual fd
  42   sockinfo
  43           calls getsockname/getpeername on the most recent
  44           <there-fd>
  45
  46   -p|-pid <pid>
  47           now attach to <pid>, detaching from previous pid
  48 '''
  49
  50 pending = []
  51 # list of (nominal, there) where nominal might be None
  52
  53 fdmap = { }
  54 # fdmap[nominal] = (actual, Donor, there)
  55
  56 def implement_pending():
  57     try: actuals = donor.fish([pend[1] for pend in pending])
  58     except fishdescriptor.fish.Error as e:
  59         print('fishdescriptor error: %s' % e, file=sys.stderr)
  60         sys.exit(127)
  61     assert(len(actuals) == len(pending))
  62     for (nominal, there), actual in zip(pending, actuals):
  63         overwriting_info = fdmap.get(nominal)
  64         if overwriting_info is not None: os.close(overwriting_info[0])
  65         fdmap[nominal] = (actual, donor, there)
  66
  67 def implement_sockinfo(nominal):
  68     (actual, tdonor, there) = fdmap[nominal]
  69     # socket.fromfd requires the AF.  But of course we don't know the AF.
  70     # There isn't a sane way to get it in Python:
  71     #  https://utcc.utoronto.ca/~cks/space/blog/python/SocketFromFdMistake
  72     # Rejected options:
  73     #  https://github.com/tiran/socketfromfd
  74     #   adds a dependency, not portable due to reliance on SO_DOMAIN
  75     #  call getsockname using ctypes
  76     #   no sane way to discover how to unpack sa_family_t
  77     perl_script = '''
  78         use strict;
  79         use Socket;
  80         use POSIX;
  81         my $sa = getsockname STDIN;
  82         exit 0 if !defined $sa and $!==ENOTSOCK;
  83         my $family = sockaddr_family $sa;
  84         print $family, "\n" or die $!;
  85     '''
  86     famp = subprocess.Popen(
  87         stdin = actual,
  88         stdout = subprocess.PIPE,
  89         args = ['perl','-we',perl_script]
  90     )
  91     (output, dummy) = famp.communicate()
  92     family = int(output)
  93
  94     sock = socket.fromfd(actual, family, 0)
  95
  96     print("[%s] %d sockinfo" % (tdonor.pid, there), end='')
  97     for f in (lambda: socket.AddressFamily(family).name,
  98               lambda: repr(sock.getsockname()),
  99               lambda: repr(sock.getpeername())):
 100         try: info = f()
 101         except Exception as e: info = repr(e)
 102         print("\t", info, sep='', end='')
 103     print("")
 104
 105     sock.close()
 106
 107 def permute_fds_for_exec():
 108     actual2intended = { info[0]: nominal for nominal, info in fdmap.items() }
 109     # invariant at the start of each loop iteration:
 110     #     for each intended (aka `nominal') we have processed:
 111     #         relevant open-file is only held in fd intended
 112     #         (unless `nominal' is None in which case it is closed)
 113     #     for each intended (aka `nominal') we have NOT processed:
 114     #         relevant open-file is only held in actual
 115     #         where  actual = fdmap[nominal][0]
 116     #         and where  actual2intended[actual] = nominal
 117     # we can rely on processing each intended only once,
 118     #  since they're hash keys
 119     # the post-condition is not really a valid state (fdmap
 120     #  is nonsense) but we call this function just before exec
 121     for intended, (actual, tdonor, there) in fdmap.items():
 122         if intended == actual:
 123             continue
 124         if intended is not None:
 125             inway_intended = actual2intended.get(intended)
 126             if inway_intended is not None:
 127                 inway_moved = os.dup(intended)
 128                 actual2intended[inway_moved] = inway_intended
 129                 fdmap[inway_intented][0] = inway_moved
 130             os.dup2(actual, intended)
 131         os.close(actual)
 132
 133 def implement_exec(argl):
 134     if donor is not None: donor.detach()
 135     sys.stdout.flush()
 136     permute_fds_for_exec()
 137     os.execvp(argl[0], argl)
 138
 139 def set_donor(pid):
 140     global donor
 141     if donor is not None: donor.detach()
 142     donor = fishdescriptor.fish.Donor(pid, debug=ov.debug)
 143
 144 def ocb_set_donor(option, opt, value, parser):
 145     set_donor(value)
 146
 147 ov = optparse.Values()
 148
 149 def process_args():
 150     global ov
 151
 152     m = None
 153
 154     def arg_matches(regexp):
 155         nonlocal m
 156         m = re.search(regexp, arg)
 157         return m
 158
 159     op = optparse.OptionParser(usage=usage)
 160
 161     op.disable_interspersed_args()
 162     op.add_option('-p','--pid', type='int', action='callback',
 163                   callback=ocb_set_donor)
 164     op.add_option('-D','--debug', action='store_const',
 165                   dest='debug', const=sys.stderr)
 166     ov.debug = None
 167
 168     args = sys.argv[1:]
 169     last_nominal = None # None or (nominal,) ie None or (None,) or (int,)
 170
 171     while True:
 172         (ov, args) = op.parse_args(args=args, values=ov)
 173         if not len(args): break
 174
 175         arg = args.pop(0)
 176
 177         if donor is None:
 178             set_donor(int(arg))
 179         elif arg_matches(r'^(?:(\d+)=)?(\d+)?$'):
 180             (nominal, there) = m.groups()
 181             nominal = None if nominal is None else int(nominal)
 182             there = int(there)
 183             pending.append((nominal,there))
 184             last_nominal = (nominal,)
 185         elif arg == 'exec':
 186             if not len(args):
 187                 op.error("exec needs command to run")
 188             implement_pending()
 189             implement_exec(args)
 190         elif arg == 'sockinfo':
 191             if last_nominal is None:
 192                 op.error('sockinfo needs a prior fd spec')
 193             implement_pending()
 194             implement_sockinfo(last_nominal[0])
 195         else:
 196             op.error("unknown argument/option `%s'" % arg)
 197
 198 process_args()