examples/python/auth/customized_auth_server.py

   1 # Copyright 2019 The gRPC Authors
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #     http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 """Server of the Python example of customizing authentication mechanism."""
  15
  16 from __future__ import absolute_import
  17 from __future__ import division
  18 from __future__ import print_function
  19
  20 import argparse
  21 import contextlib
  22 import logging
  23 from concurrent import futures
  24
  25 import grpc
  26 from examples import helloworld_pb2
  27 from examples import helloworld_pb2_grpc
  28 from examples.python.auth import _credentials
  29
  30 _LOGGER = logging.getLogger(__name__)
  31 _LOGGER.setLevel(logging.INFO)
  32
  33 _LISTEN_ADDRESS_TEMPLATE = 'localhost:%d'
  34 _SIGNATURE_HEADER_KEY = 'x-signature'
  35
  36
  37 class SignatureValidationInterceptor(grpc.ServerInterceptor):
  38
  39     def __init__(self):
  40
  41         def abort(ignored_request, context):
  42             context.abort(grpc.StatusCode.UNAUTHENTICATED, 'Invalid signature')
  43
  44         self._abortion = grpc.unary_unary_rpc_method_handler(abort)
  45
  46     def intercept_service(self, continuation, handler_call_details):
  47         # Example HandlerCallDetails object:
  48         #     _HandlerCallDetails(
  49         #       method=u'/helloworld.Greeter/SayHello',
  50         #       invocation_metadata=...)
  51         method_name = handler_call_details.method.split('/')[-1]
  52         expected_metadata = (_SIGNATURE_HEADER_KEY, method_name[::-1])
  53         if expected_metadata in handler_call_details.invocation_metadata:
  54             return continuation(handler_call_details)
  55         else:
  56             return self._abortion
  57
  58
  59 class SimpleGreeter(helloworld_pb2_grpc.GreeterServicer):
  60
  61     def SayHello(self, request, unused_context):
  62         return helloworld_pb2.HelloReply(message='Hello, %s!' % request.name)
  63
  64
  65 @contextlib.contextmanager
  66 def run_server(port):
  67     # Bind interceptor to server
  68     server = grpc.server(futures.ThreadPoolExecutor(),
  69                          interceptors=(SignatureValidationInterceptor(),))
  70     helloworld_pb2_grpc.add_GreeterServicer_to_server(SimpleGreeter(), server)
  71
  72     # Loading credentials
  73     server_credentials = grpc.ssl_server_credentials(((
  74         _credentials.SERVER_CERTIFICATE_KEY,
  75         _credentials.SERVER_CERTIFICATE,
  76     ),))
  77
  78     # Pass down credentials
  79     port = server.add_secure_port(_LISTEN_ADDRESS_TEMPLATE % port,
  80                                   server_credentials)
  81
  82     server.start()
  83     try:
  84         yield server, port
  85     finally:
  86         server.stop(0)
  87
  88
  89 def main():
  90     parser = argparse.ArgumentParser()
  91     parser.add_argument('--port',
  92                         nargs='?',
  93                         type=int,
  94                         default=50051,
  95                         help='the listening port')
  96     args = parser.parse_args()
  97
  98     with run_server(args.port) as (server, port):
  99         logging.info('Server is listening at port :%d', port)
 100         server.wait_for_termination()
 101
 102
 103 if __name__ == '__main__':
 104     logging.basicConfig(level=logging.INFO)
 105     main()