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 16import argparse 17from concurrent import futures 18import contextlib 19import logging 20 21import _credentials 22import grpc 23 24helloworld_pb2, helloworld_pb2_grpc = grpc.protos_and_services( 25 "helloworld.proto" 26) 27 28_LOGGER = logging.getLogger(__name__) 29_LOGGER.setLevel(logging.INFO) 30 31_LISTEN_ADDRESS_TEMPLATE = "localhost:%d" 32_SIGNATURE_HEADER_KEY = "x-signature" 33 34 35class SignatureValidationInterceptor(grpc.ServerInterceptor): 36 def __init__(self): 37 def abort(ignored_request, context): 38 context.abort(grpc.StatusCode.UNAUTHENTICATED, "Invalid signature") 39 40 self._abortion = grpc.unary_unary_rpc_method_handler(abort) 41 42 def intercept_service(self, continuation, handler_call_details): 43 # Example HandlerCallDetails object: 44 # _HandlerCallDetails( 45 # method=u'/helloworld.Greeter/SayHello', 46 # invocation_metadata=...) 47 method_name = handler_call_details.method.split("/")[-1] 48 expected_metadata = (_SIGNATURE_HEADER_KEY, method_name[::-1]) 49 if expected_metadata in handler_call_details.invocation_metadata: 50 return continuation(handler_call_details) 51 else: 52 return self._abortion 53 54 55class SimpleGreeter(helloworld_pb2_grpc.GreeterServicer): 56 def SayHello(self, request, unused_context): 57 return helloworld_pb2.HelloReply(message="Hello, %s!" % request.name) 58 59 60@contextlib.contextmanager 61def run_server(port): 62 # Bind interceptor to server 63 server = grpc.server( 64 futures.ThreadPoolExecutor(), 65 interceptors=(SignatureValidationInterceptor(),), 66 ) 67 helloworld_pb2_grpc.add_GreeterServicer_to_server(SimpleGreeter(), server) 68 69 # Loading credentials 70 server_credentials = grpc.ssl_server_credentials( 71 ( 72 ( 73 _credentials.SERVER_CERTIFICATE_KEY, 74 _credentials.SERVER_CERTIFICATE, 75 ), 76 ) 77 ) 78 79 # Pass down credentials 80 port = server.add_secure_port( 81 _LISTEN_ADDRESS_TEMPLATE % port, server_credentials 82 ) 83 84 server.start() 85 try: 86 yield server, port 87 finally: 88 server.stop(0) 89 90 91def main(): 92 parser = argparse.ArgumentParser() 93 parser.add_argument( 94 "--port", nargs="?", type=int, default=50051, help="the listening port" 95 ) 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 103if __name__ == "__main__": 104 logging.basicConfig(level=logging.INFO) 105 main() 106