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