numam-spdk/scripts/rpc_http_proxy.py

#!/usr/bin/env python3

import base64
import errno
import json
import socket
import sys
try:
    from BaseHTTPServer import HTTPServer, BaseHTTPRequestHandler
except ImportError:
    from http.server import HTTPServer
    from http.server import BaseHTTPRequestHandler

rpc_sock = None


def print_usage_and_exit(status):
    print('Usage: rpc_http_proxy.py <server IP> <server port> <user name>' +
          ' <password> <SPDK RPC socket (optional, default: /var/tmp/spdk.sock)>')
    sys.exit(status)


def rpc_call(req):
    global rpc_sock

    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.connect(rpc_sock)
    sock.sendall(req)

    if 'id' not in json.loads(req.decode('ascii')):
        sock.close()
        return None

    buf = ''
    closed = False
    response = None

    while not closed:
        newdata = sock.recv(1024)
        if (newdata == b''):
            closed = True
        buf += newdata.decode('ascii')
        try:
            response = json.loads(buf)
        except ValueError:
            continue  # incomplete response; keep buffering
        break

    sock.close()

    if not response and len(buf) > 0:
        raise

    return buf


class ServerHandler(BaseHTTPRequestHandler):

    key = ""

    def do_HEAD(self):
        self.send_response(200)
        self.send_header('Content-type', 'text/html')
        self.end_headers()

    def do_AUTHHEAD(self):
        self.send_response(401)
        self.send_header('WWW-Authenticate', 'text/html')
        self.send_header('Content-type', 'text/html')
        self.end_headers()

    def do_INTERNALERROR(self):
        self.send_response(500)
        self.send_header('Content-type', 'text/html')
        self.end_headers()

    def do_POST(self):
        if self.headers['Authorization'] != 'Basic ' + self.key:
            self.do_AUTHHEAD()
        else:
            data_string = self.rfile.read(int(self.headers['Content-Length']))

            try:
                response = rpc_call(data_string)
                if response is not None:
                    self.do_HEAD()
                    self.wfile.write(bytes(response.encode(encoding='ascii')))
            except ValueError:
                self.do_INTERNALERROR()


def main():
    global rpc_sock

    if len(sys.argv) == 1 or sys.argv[1] == '-h':
        print_usage_and_exit(0)
    elif len(sys.argv) < 5:
        print('Not enough arguments')
        print_usage_and_exit(errno.EINVAL)
    elif len(sys.argv) > 6:
        print('Too many arguments')
        print_usage_and_exit(errno.E2BIG)

    if len(sys.argv) == 6:
        rpc_sock = sys.argv[5]
    else:
        rpc_sock = '/var/tmp/spdk.sock'

    # encoding user name and password
    key = base64.b64encode((sys.argv[3]+':'+sys.argv[4]).encode(encoding='ascii')).decode('ascii')

    try:
        ServerHandler.key = key
        httpd = HTTPServer((sys.argv[1], int(sys.argv[2])), ServerHandler)
        print('Started RPC http proxy server')
        httpd.serve_forever()
    except KeyboardInterrupt:
        print('Shutting down server')
        httpd.socket.close()


if __name__ == '__main__':
    main()