| 1 | #!/usr/bin/env python2
 | 
| 2 | """
 | 
| 3 | Utility for checking the output of group-session.sh
 | 
| 4 | """
 | 
| 5 | from __future__ import print_function
 | 
| 6 | 
 | 
| 7 | import re
 | 
| 8 | import sys
 | 
| 9 | 
 | 
| 10 | 
 | 
| 11 | class Process(object):
 | 
| 12 | 
 | 
| 13 |   def __init__(self, pid, ppid, pgid, comm):
 | 
| 14 |     self.pid = pid
 | 
| 15 |     self.ppid = ppid
 | 
| 16 |     self.pgid = pgid
 | 
| 17 |     self.comm = comm
 | 
| 18 | 
 | 
| 19 |   def __str__(self):
 | 
| 20 |     return '\t'.join((self.pid, self.ppid, self.pgid, self.comm))
 | 
| 21 | 
 | 
| 22 |   def assert_pgid(self, pgid):
 | 
| 23 |     if self.pgid != pgid:
 | 
| 24 |       print('[%s] has pgid %s. expected %s.' %
 | 
| 25 |           (self, self.pgid, pgid), file=sys.stderr)
 | 
| 26 |       sys.exit(1)
 | 
| 27 | 
 | 
| 28 | class ProcessTree(object):
 | 
| 29 | 
 | 
| 30 |   def __init__(self, proc):
 | 
| 31 |     self.proc = proc
 | 
| 32 |     self.children = []
 | 
| 33 |   
 | 
| 34 |   def __str__(self):
 | 
| 35 |     lines = [str(self.proc)]
 | 
| 36 |     for child in self.children:
 | 
| 37 |       lines.append(str(child))
 | 
| 38 | 
 | 
| 39 |     return '\n'.join(lines)
 | 
| 40 | 
 | 
| 41 |   def assert_child_count(self, n):
 | 
| 42 |     if len(self.children) != n:
 | 
| 43 |       print('[%s] has %d children. expected %d.' %
 | 
| 44 |           (self.proc, len(self.children), n), file=sys.stderr)
 | 
| 45 |       sys.exit(1)
 | 
| 46 | 
 | 
| 47 | 
 | 
| 48 | def parse_process_tree(f, runner_pid):
 | 
| 49 |   procs = {}
 | 
| 50 | 
 | 
| 51 |   for line in f:
 | 
| 52 |     m = re.match(r'^\s*(\d+)\s+(\d+)\s+(\d+)\s+(\d+)\s+(-?\d+)\s+(\w+)$', line)
 | 
| 53 |     if not m:
 | 
| 54 |       continue
 | 
| 55 |     # TODO: use SID and TPGID
 | 
| 56 |     pid, ppid, pgid, _, _, comm = m.groups()
 | 
| 57 |     proc = Process(pid, ppid, pgid, comm)
 | 
| 58 |     ptree = ProcessTree(proc)
 | 
| 59 |     procs[proc.pid] = ptree
 | 
| 60 |     if proc.ppid in procs:
 | 
| 61 |       procs[proc.ppid].children.append(ptree)
 | 
| 62 | 
 | 
| 63 |   if runner_pid not in procs:
 | 
| 64 |     print('malformed ps output', file=sys.stderr)
 | 
| 65 |     sys.exit(1)
 | 
| 66 | 
 | 
| 67 |   # first process is the test harness
 | 
| 68 |   root = procs[runner_pid]
 | 
| 69 |   root.assert_child_count(1)
 | 
| 70 |   return root.children[0]
 | 
| 71 | 
 | 
| 72 | 
 | 
| 73 | def check_proc(ptree, shell, interactive):
 | 
| 74 |   assert len(ptree.children) == 1
 | 
| 75 |   ps = ptree.children[0]
 | 
| 76 |   if interactive:
 | 
| 77 |     ps.proc.assert_pgid(ps.proc.pid)
 | 
| 78 |   else:
 | 
| 79 |     ps.proc.assert_pgid(ptree.proc.pgid)
 | 
| 80 | 
 | 
| 81 | 
 | 
| 82 | def check_pipe(ptree, shell, snippet, interactive):
 | 
| 83 |   if snippet == 'fgpipe-lastpipe' and ('zsh' in shell or 'osh' in shell):
 | 
| 84 |     expected_children = 2
 | 
| 85 |   else:
 | 
| 86 |     expected_children = 3
 | 
| 87 | 
 | 
| 88 |   ptree.assert_child_count(expected_children)
 | 
| 89 | 
 | 
| 90 |   first = None
 | 
| 91 |   for child in ptree.children:
 | 
| 92 |     if child.proc.pid == child.proc.pgid:
 | 
| 93 |       first = child
 | 
| 94 |       break
 | 
| 95 | 
 | 
| 96 |   if not first and interactive:
 | 
| 97 |     print('interactive pipeline has no leader', file=sys.stderr) 
 | 
| 98 |     sys.exit(1)
 | 
| 99 | 
 | 
| 100 |   pgid = first.proc.pgid if first else ptree.proc.pgid
 | 
| 101 | 
 | 
| 102 |   for child in ptree.children:
 | 
| 103 |     child.proc.assert_pgid(pgid)
 | 
| 104 | 
 | 
| 105 | 
 | 
| 106 | def check_subshell(ptree, shell, interactive):
 | 
| 107 |   ptree.assert_child_count(1)
 | 
| 108 |   subshell = ptree.children[0]
 | 
| 109 |   subshell.assert_child_count(1)
 | 
| 110 |   ps = subshell.children[0]
 | 
| 111 | 
 | 
| 112 |   if interactive:
 | 
| 113 |     subshell.proc.assert_pgid(subshell.proc.pid)
 | 
| 114 |     ps.proc.assert_pgid(subshell.proc.pid)
 | 
| 115 |   else:
 | 
| 116 |     subshell.proc.assert_pgid(ptree.proc.pgid)
 | 
| 117 |     ps.proc.assert_pgid(ptree.proc.pgid)
 | 
| 118 | 
 | 
| 119 | 
 | 
| 120 | def check_csub(ptree, shell, interactive):
 | 
| 121 |   ptree.assert_child_count(1)
 | 
| 122 |   ps = ptree.children[0]
 | 
| 123 |   ps.proc.assert_pgid(ptree.proc.pgid)
 | 
| 124 | 
 | 
| 125 | 
 | 
| 126 | def check_psub(ptree, shell, interactive):
 | 
| 127 |   ps, cat, subshell = None, None, None
 | 
| 128 |   if shell == 'bash':
 | 
| 129 |     ptree.assert_child_count(2)
 | 
| 130 |     for child in ptree.children:
 | 
| 131 |       if len(child.children) == 1:
 | 
| 132 |         subshell = child
 | 
| 133 |         ps = child.children[0]
 | 
| 134 |       elif len(child.children) == 0:
 | 
| 135 |         cat = child
 | 
| 136 |       else:
 | 
| 137 |         print('[%s] has unexpected child [%s]' % (ptree.proc, child), file=sys.stderr)
 | 
| 138 |         sys.exit(1)
 | 
| 139 | 
 | 
| 140 |     if not subshell:
 | 
| 141 |       print('missing expected subshell', file=sys.stderr)
 | 
| 142 |       sys.exit(1)
 | 
| 143 |   else:
 | 
| 144 |     ptree.assert_child_count(2)
 | 
| 145 |     # NOTE: Ideally we would check the comm field of the children, but `ps` may
 | 
| 146 |     # have run before some of them called exec(). Luckily we're only checkign
 | 
| 147 |     # that both children are in their own group in this case, so we just
 | 
| 148 |     # guess...
 | 
| 149 |     ps = ptree.children[0]
 | 
| 150 |     cat = ptree.children[1]
 | 
| 151 | 
 | 
| 152 |   if not ps:
 | 
| 153 |     print('missing ps', file=sys.stderr)
 | 
| 154 |     sys.exit(1)
 | 
| 155 | 
 | 
| 156 |   if not cat:
 | 
| 157 |     print('missing cat', file=sys.stderr)
 | 
| 158 |     sys.exit(1)
 | 
| 159 |   
 | 
| 160 | 
 | 
| 161 |   if not interactive:
 | 
| 162 |     ps.proc.assert_pgid(ptree.proc.pgid)
 | 
| 163 |     cat.proc.assert_pgid(ptree.proc.pgid)
 | 
| 164 |     if subshell:
 | 
| 165 |       subshell.proc.assert_pgid(ptree.proc.pgid)
 | 
| 166 |   else:
 | 
| 167 |     if shell == 'bash':
 | 
| 168 |       # bash is interesting
 | 
| 169 |       subshell.proc.assert_pgid(ptree.proc.pid)
 | 
| 170 |       ps.proc.assert_pgid(ptree.proc.pid)
 | 
| 171 |       cat.proc.assert_pgid(cat.proc.pid)
 | 
| 172 |     else:
 | 
| 173 |       # osh and zsh put all children in their own group
 | 
| 174 |       ps.proc.assert_pgid(ps.proc.pid)
 | 
| 175 |       cat.proc.assert_pgid(cat.proc.pid)
 | 
| 176 | 
 | 
| 177 | 
 | 
| 178 | def main(argv):
 | 
| 179 |   runner_pid = argv[1]
 | 
| 180 |   shell = argv[2]
 | 
| 181 |   snippet = argv[3]
 | 
| 182 |   interactive = (argv[4] == 'yes')
 | 
| 183 | 
 | 
| 184 |   ptree = parse_process_tree(sys.stdin, runner_pid)
 | 
| 185 |   if snippet == 'fgproc':
 | 
| 186 |     check_proc(ptree, shell, interactive)
 | 
| 187 | 
 | 
| 188 |   elif snippet == 'bgproc':
 | 
| 189 |     check_proc(ptree, shell, interactive)
 | 
| 190 | 
 | 
| 191 |   elif snippet == 'fgpipe':
 | 
| 192 |     check_pipe(ptree, shell, snippet, interactive)
 | 
| 193 | 
 | 
| 194 |   elif snippet == 'fgpipe-lastpipe':
 | 
| 195 |     check_pipe(ptree, shell, snippet, interactive)
 | 
| 196 | 
 | 
| 197 |   elif snippet == 'bgpipe':
 | 
| 198 |     check_pipe(ptree, shell, snippet, interactive)
 | 
| 199 | 
 | 
| 200 |   elif snippet == 'bgpipe-lastpipe':
 | 
| 201 |     check_pipe(ptree, shell, snippet, interactive)
 | 
| 202 | 
 | 
| 203 |   elif snippet == 'subshell':
 | 
| 204 |     check_subshell(ptree, shell, interactive)
 | 
| 205 | 
 | 
| 206 |   elif snippet == 'csub':
 | 
| 207 |     check_csub(ptree, shell, interactive)
 | 
| 208 | 
 | 
| 209 |   elif snippet == 'psub':
 | 
| 210 |     check_psub(ptree, shell, interactive)
 | 
| 211 | 
 | 
| 212 |   else:
 | 
| 213 |     raise RuntimeError('Invalid snippet %r' % snippet)
 | 
| 214 | 
 | 
| 215 |   return 0
 | 
| 216 | 
 | 
| 217 | 
 | 
| 218 | if __name__ == '__main__':
 | 
| 219 |   try:
 | 
| 220 |     sys.exit(main(sys.argv))
 | 
| 221 |   except RuntimeError as e:
 | 
| 222 |     print('FATAL: %s' % e, file=sys.stderr)
 | 
| 223 |     sys.exit(1)
 |