1#!/usr/bin/env python3
2#
3#   Copyright 2020 - The Android Open Source Project
4#
5#   Licensed under the Apache License, Version 2.0 (the "License");
6#   you may not use this file except in compliance with the License.
7#   You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0
10#
11#   Unless required by applicable law or agreed to in writing, software
12#   distributed under the License is distributed on an "AS IS" BASIS,
13#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#   See the License for the specific language governing permissions and
15#   limitations under the License.
16
17import grpc
18import re
19
20from facade import common_pb2 as common
21from google.protobuf import text_format
22
23
24def custom_message_formatter(m, ident, as_one_line):
25    if m.DESCRIPTOR == common.Data.DESCRIPTOR:
26        return 'payload: (hex) "{}"'.format(m.payload.hex(" "))
27    return None
28
29
30def pretty_print(request):
31    return '{} {}'.format(
32        type(request).__name__,
33        text_format.MessageToString(request, as_one_line=True, message_formatter=custom_message_formatter))
34
35
36class LoggingRandezvousWrapper():
37
38    def __init__(self, server_stream_call, logTag):
39        if server_stream_call is None:
40            raise ValueError("server_stream_call cannot be None")
41        self.server_stream_call = server_stream_call
42        self.logTag = logTag
43
44    def cancel(self):
45        self.server_stream_call.cancel()
46
47    def cancelled(self):
48        return self.server_stream_call.cancelled()
49
50    def __iter__(self):
51        return self
52
53    def __next__(self):
54        resp = self.server_stream_call.__next__()
55        print("%s %s" % (self.logTag, pretty_print(resp)))
56        return resp
57
58
59class LoggingClientInterceptor(grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor):
60
61    TAG_MIN_WIDTH = 24
62
63    def __init__(self, name):
64        self.name = name
65        self.inLogTag = "[host ▶▶▶▶▶ %s]" % self.name
66        self.outLogTag = "[host ◀◀◀◀◀ %s]" % self.name
67        tagLength = len(re.sub('[^\w\s]', '', self.inLogTag)) + 11
68        if tagLength < self.TAG_MIN_WIDTH:
69            self.inLogTag += " " * (self.TAG_MIN_WIDTH - tagLength)
70            self.outLogTag += " " * (self.TAG_MIN_WIDTH - tagLength)
71
72    def intercept_unary_unary(self, continuation, client_call_details, request):
73        """
74        This interceptor logs the requests from host
75        """
76        print("%s%s %s" % (self.inLogTag, client_call_details.method, pretty_print(request)))
77        return continuation(client_call_details, request)
78
79    def intercept_unary_stream(self, continuation, client_call_details, request):
80        """
81        This interceptor wraps the server response, and logs all the messages coming to host
82        """
83        print("%s%s %s" % (self.inLogTag, client_call_details.method, pretty_print(request)))
84        server_stream_call = continuation(client_call_details, request)
85        retuningMsgLogTag = self.outLogTag + client_call_details.method
86        return LoggingRandezvousWrapper(server_stream_call, retuningMsgLogTag)
87