| 1 | """
 | 
| 2 | py_fanos.py: Pure Python implementation of FANOS
 | 
| 3 | 
 | 
| 4 | Python 2 doesn't have native FD passing, but Python 3 does.
 | 
| 5 | """
 | 
| 6 | 
 | 
| 7 | import array
 | 
| 8 | import os
 | 
| 9 | import socket
 | 
| 10 | import sys
 | 
| 11 | 
 | 
| 12 | 
 | 
| 13 | ARGV0 = os.path.basename(sys.argv[0])
 | 
| 14 | 
 | 
| 15 | def log(msg, *args):
 | 
| 16 |   if args:
 | 
| 17 |     msg = msg % args
 | 
| 18 |   print('[%s] %s' %(ARGV0, msg), file=sys.stderr)
 | 
| 19 | 
 | 
| 20 | 
 | 
| 21 | def send(sock, msg, fds=None):
 | 
| 22 |   """Send a blob and optional file descriptors."""
 | 
| 23 | 
 | 
| 24 |   fds = fds or []
 | 
| 25 | 
 | 
| 26 |   sock.send(b'%d:' % len(msg))  # netstring prefix
 | 
| 27 | 
 | 
| 28 |   # Send the FILE DESCRIPTOR with the NETSTRING PAYLOAD
 | 
| 29 |   ancillary = (
 | 
| 30 |     socket.SOL_SOCKET, socket.SCM_RIGHTS, array.array("i", fds)
 | 
| 31 |   )
 | 
| 32 |   result = sock.sendmsg([msg], [ancillary])
 | 
| 33 |   #log('sendmsg returned %s', result)
 | 
| 34 | 
 | 
| 35 |   sock.send(b',')  # trailing netstring thing
 | 
| 36 | 
 | 
| 37 | 
 | 
| 38 | def recv_fds_once(sock, msglen, maxfds, fd_out):
 | 
| 39 |   """Helper function from Python stdlib docs."""
 | 
| 40 |   fds = array.array("i")   # Array of ints
 | 
| 41 |   msg, ancdata, flags, addr = sock.recvmsg(msglen,
 | 
| 42 |                                            socket.CMSG_LEN(maxfds * fds.itemsize))
 | 
| 43 |   for cmsg_level, cmsg_type, cmsg_data in ancdata:
 | 
| 44 |     if cmsg_level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS:
 | 
| 45 |       # Append data, ignoring any truncated integers at the end.
 | 
| 46 |       fds.frombytes(cmsg_data[:len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
 | 
| 47 | 
 | 
| 48 |   fd_out.extend(fds)
 | 
| 49 |   return msg
 | 
| 50 | 
 | 
| 51 | 
 | 
| 52 | def recv(sock, fd_out=None):
 | 
| 53 |   """Receive a blob and optional file descriptors.
 | 
| 54 | 
 | 
| 55 |   Returns:
 | 
| 56 |     The message blob, or None when the other end closes at a valid message
 | 
| 57 |     boundary.
 | 
| 58 | 
 | 
| 59 |     Appends to fd_out.
 | 
| 60 |   """
 | 
| 61 |   if fd_out is None:
 | 
| 62 |     fd_out = []  # Can be thrown away
 | 
| 63 | 
 | 
| 64 |   len_buf = []
 | 
| 65 |   for i in range(10):
 | 
| 66 |     byte = sock.recv(1)
 | 
| 67 |     #log('byte = %r', byte)
 | 
| 68 | 
 | 
| 69 |     # This happens on close()
 | 
| 70 |     if len(byte) == 0:
 | 
| 71 |       if i == 0:
 | 
| 72 |         return None  # that was the last message
 | 
| 73 |       else:
 | 
| 74 |         raise ValueError('Unexpected EOF')
 | 
| 75 | 
 | 
| 76 |     if b'0' <= byte and byte <= b'9':
 | 
| 77 |       len_buf.append(byte)
 | 
| 78 |     else:
 | 
| 79 |       break
 | 
| 80 | 
 | 
| 81 |   if len(len_buf) == 0:
 | 
| 82 |     raise ValueError('Expected netstring length')
 | 
| 83 |   if byte != b':':
 | 
| 84 |     raise ValueError('Expected : after length')
 | 
| 85 | 
 | 
| 86 |   num_bytes = int(b''.join(len_buf))
 | 
| 87 |   #log('num_bytes = %d', num_bytes)
 | 
| 88 | 
 | 
| 89 |   msg = b''
 | 
| 90 |   while True:
 | 
| 91 |     chunk = recv_fds_once(sock, num_bytes, 3, fd_out)
 | 
| 92 |     #log("chunk %r  FDs %s", chunk, fds)
 | 
| 93 | 
 | 
| 94 |     msg += chunk
 | 
| 95 |     if len(msg) == num_bytes:
 | 
| 96 |       break
 | 
| 97 | 
 | 
| 98 |   byte = sock.recv(1)
 | 
| 99 |   if byte != b',':
 | 
| 100 |     raise ValueError('Expected ,')
 | 
| 101 | 
 | 
| 102 |   return msg
 |