xref: /aosp_15_r20/external/grpc-grpc/examples/python/auth/customized_auth_client.py (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2)
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"""Client of the Python example of customizing authentication mechanism."""
15
16import argparse
17import contextlib
18import logging
19
20import _credentials
21import grpc
22
23helloworld_pb2, helloworld_pb2_grpc = grpc.protos_and_services(
24    "helloworld.proto"
25)
26
27_LOGGER = logging.getLogger(__name__)
28_LOGGER.setLevel(logging.INFO)
29
30_SERVER_ADDR_TEMPLATE = "localhost:%d"
31_SIGNATURE_HEADER_KEY = "x-signature"
32
33
34class AuthGateway(grpc.AuthMetadataPlugin):
35    def __call__(self, context, callback):
36        """Implements authentication by passing metadata to a callback.
37
38        Implementations of this method must not block.
39
40        Args:
41          context: An AuthMetadataContext providing information on the RPC that
42            the plugin is being called to authenticate.
43          callback: An AuthMetadataPluginCallback to be invoked either
44            synchronously or asynchronously.
45        """
46        # Example AuthMetadataContext object:
47        # AuthMetadataContext(
48        #     service_url=u'https://localhost:50051/helloworld.Greeter',
49        #     method_name=u'SayHello')
50        signature = context.method_name[::-1]
51        # NOTE: The metadata keys provided to the callback must be lower-cased.
52        callback(((_SIGNATURE_HEADER_KEY, signature),), None)
53
54
55@contextlib.contextmanager
56def create_client_channel(addr):
57    # Call credential object will be invoked for every single RPC
58    call_credentials = grpc.metadata_call_credentials(
59        AuthGateway(), name="auth gateway"
60    )
61    # Channel credential will be valid for the entire channel
62    channel_credential = grpc.ssl_channel_credentials(
63        _credentials.ROOT_CERTIFICATE
64    )
65    # Combining channel credentials and call credentials together
66    composite_credentials = grpc.composite_channel_credentials(
67        channel_credential,
68        call_credentials,
69    )
70    channel = grpc.secure_channel(addr, composite_credentials)
71    yield channel
72
73
74def send_rpc(channel):
75    stub = helloworld_pb2_grpc.GreeterStub(channel)
76    request = helloworld_pb2.HelloRequest(name="you")
77    try:
78        response = stub.SayHello(request)
79    except grpc.RpcError as rpc_error:
80        _LOGGER.error("Received error: %s", rpc_error)
81        return rpc_error
82    else:
83        _LOGGER.info("Received message: %s", response)
84        return response
85
86
87def main():
88    parser = argparse.ArgumentParser()
89    parser.add_argument(
90        "--port",
91        nargs="?",
92        type=int,
93        default=50051,
94        help="the address of server",
95    )
96    args = parser.parse_args()
97
98    with create_client_channel(_SERVER_ADDR_TEMPLATE % args.port) as channel:
99        send_rpc(channel)
100
101
102if __name__ == "__main__":
103    logging.basicConfig(level=logging.INFO)
104    main()
105