qubesd_query.py 2.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. #!/usr/bin/env python3.6
  2. import argparse
  3. import asyncio
  4. import signal
  5. import sys
  6. QUBESD_SOCK = '/var/run/qubesd.sock'
  7. parser = argparse.ArgumentParser(
  8. description='low-level qubesd interrogation tool')
  9. parser.add_argument('--connect', '-c', metavar='PATH',
  10. dest='socket',
  11. default=QUBESD_SOCK,
  12. help='path to qubesd UNIX socket (default: %(default)s)')
  13. parser.add_argument('--empty', '-e',
  14. dest='payload',
  15. action='store_false', default=True,
  16. help='do not read from stdin and send empty payload')
  17. parser.add_argument('--fail',
  18. dest='fail',
  19. action='store_true',
  20. help='Should non-OK qubesd response result in non-zero exit code')
  21. parser.add_argument('src', metavar='SRC',
  22. help='source qube')
  23. parser.add_argument('method', metavar='METHOD',
  24. help='method name')
  25. parser.add_argument('dest', metavar='DEST',
  26. help='destination qube')
  27. parser.add_argument('arg', metavar='ARGUMENT',
  28. nargs='?', default='',
  29. help='argument to method')
  30. def sighandler(loop, signame, coro):
  31. print('caught {}, exiting'.format(signame))
  32. coro.cancel()
  33. loop.stop()
  34. @asyncio.coroutine
  35. def qubesd_client(socket, payload, *args):
  36. '''
  37. Connect to qubesd, send request and passthrough response to stdout
  38. :param socket: path to qubesd socket
  39. :param payload: payload of the request
  40. :param args: request to qubesd
  41. :return:
  42. '''
  43. try:
  44. reader, writer = yield from asyncio.open_unix_connection(socket)
  45. except asyncio.CancelledError:
  46. return 1
  47. for arg in args:
  48. writer.write(arg.encode('ascii'))
  49. writer.write(b'\0')
  50. writer.write(payload)
  51. writer.write_eof()
  52. try:
  53. header_data = yield from reader.read(1)
  54. returncode = int(header_data)
  55. sys.stdout.buffer.write(header_data) # pylint: disable=no-member
  56. while not reader.at_eof():
  57. data = yield from reader.read(4096)
  58. sys.stdout.buffer.write(data) # pylint: disable=no-member
  59. sys.stdout.flush()
  60. return returncode
  61. except asyncio.CancelledError:
  62. return 1
  63. finally:
  64. writer.close()
  65. def main(args=None):
  66. args = parser.parse_args(args)
  67. loop = asyncio.get_event_loop()
  68. # pylint: disable=no-member
  69. payload = sys.stdin.buffer.read() if args.payload else b''
  70. # pylint: enable=no-member
  71. coro = asyncio.ensure_future(qubesd_client(args.socket, payload,
  72. args.src, args.method, args.dest, args.arg))
  73. for signame in ('SIGINT', 'SIGTERM'):
  74. loop.add_signal_handler(getattr(signal, signame),
  75. sighandler, loop, signame, coro)
  76. try:
  77. returncode = loop.run_until_complete(coro)
  78. finally:
  79. loop.close()
  80. if args.fail:
  81. return returncode
  82. return 0
  83. if __name__ == '__main__':
  84. sys.exit(main())