text: # Copyright 2020 gRPC authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """
  This contains helpers for gRPC services defined in
  https://github.com/grpc/grpc-proto/blob/master/grpc/channelz/v1/channelz.proto
  """
  import ipaddress
  import logging
  from typing import Iterator, Optional
  
  import grpc
  from grpc_channelz.v1 import channelz_pb2
  from grpc_channelz.v1 import channelz_pb2_grpc
  
  import framework.rpc
  
  logger = logging.getLogger(__name__)
  
  # Type aliases
  # Channel
  Channel = channelz_pb2.Channel
  ChannelConnectivityState = channelz_pb2.ChannelConnectivityState
  ChannelState = ChannelConnectivityState.State  # pylint: disable=no-member
  _GetTopChannelsRequest = channelz_pb2.GetTopChannelsRequest
  _GetTopChannelsResponse = channelz_pb2.GetTopChannelsResponse
  # Subchannel
  Subchannel = channelz_pb2.Subchannel
  _GetSubchannelRequest = channelz_pb2.GetSubchannelRequest
  _GetSubchannelResponse = channelz_pb2.GetSubchannelResponse
  # Server
  Server = channelz_pb2.Server
  _GetServersRequest = channelz_pb2.GetServersRequest
  _GetServersResponse = channelz_pb2.GetServersResponse
  # Sockets
  Socket = channelz_pb2.Socket
  SocketRef = channelz_pb2.SocketRef
  _GetSocketRequest = channelz_pb2.GetSocketRequest
  _GetSocketResponse = channelz_pb2.GetSocketResponse
  Address = channelz_pb2.Address
  Security = channelz_pb2.Security
  # Server Sockets
  _GetServerSocketsRequest = channelz_pb2.GetServerSocketsRequest
  _GetServerSocketsResponse = channelz_pb2.GetServerSocketsResponse
  
  
  class ChannelzServiceClient(framework.rpc.grpc.GrpcClientHelper):
      stub: channelz_pb2_grpc.ChannelzStub
  
      def __init__(self, channel: grpc.Channel):
          super().__init__(channel, channelz_pb2_grpc.ChannelzStub)
  
      @staticmethod
      def is_sock_tcpip_address(address: Address):
          return address.WhichOneof('address') == 'tcpip_address'
  
      @staticmethod
      def is_ipv4(tcpip_address: Address.TcpIpAddress):
          # According to proto, tcpip_address.ip_address is either IPv4 or IPv6.
          # Correspondingly, it's either 4 bytes or 16 bytes in length.
          return len(tcpip_address.ip_address) == 4
  
      @classmethod
      def sock_address_to_str(cls, address: Address):
          if cls.is_sock_tcpip_address(address):
              tcpip_address: Address.TcpIpAddress = address.tcpip_address
              if cls.is_ipv4(tcpip_address):
                  ip = ipaddress.IPv4Address(tcpip_address.ip_address)
              else:
                  ip = ipaddress.IPv6Address(tcpip_address.ip_address)
              return f'{ip}:{tcpip_address.port}'
          else:
              raise NotImplementedError('Only tcpip_address implemented')
  
      @classmethod
      def sock_addresses_pretty(cls, socket: Socket):
          return (f'local={cls.sock_address_to_str(socket.local)}, '
                  f'remote={cls.sock_address_to_str(socket.remote)}')
  
      @staticmethod
      def find_server_socket_matching_client(server_sockets: Iterator[Socket],
                                             client_socket: Socket) -> Socket:
          for server_socket in server_sockets:
              if server_socket.remote == client_socket.local:
                  return server_socket
          return None
  
      def find_channels_for_target(self, target: str) -> Iterator[Channel]:
          return (channel for channel in self.list_channels()
                  if channel.data.target == target)
  
      def find_server_listening_on_port(self, port: int) -> Optional[Server]:
          for server in self.list_servers():
              listen_socket_ref: SocketRef
              for listen_socket_ref in server.listen_socket:
                  listen_socket = self.get_socket(listen_socket_ref.socket_id)
                  listen_address: Address = listen_socket.local
                  if (self.is_sock_tcpip_address(listen_address) and
                          listen_address.tcpip_address.port == port):
                      return server
          return None
  
      def list_channels(self) -> Iterator[Channel]:
          """
          Iterate over all pages of all root channels.
  
          Root channels are those which application has directly created.
          This does not include subchannels nor non-top level channels.
          """
          start: int = -1
          response: Optional[_GetTopChannelsResponse] = None
          while start < 0 or not response.end:
              # From proto: To request subsequent pages, the client generates this
              # value by adding 1 to the highest seen result ID.
              start += 1
              response = self.call_unary_with_deadline(
                  rpc='GetTopChannels',
                  req=_GetTopChannelsRequest(start_channel_id=start))
              for channel in response.channel:
                  start = max(start, channel.ref.channel_id)
                  yield channel
  
      def list_servers(self) -> Iterator[Server]:
          """Iterate over all pages of all servers that exist in the process."""
          start: int = -1
          response: Optional[_GetServersResponse] = None
          while start < 0 or not response.end:
              # From proto: To request subsequent pages, the client generates this
              # value by adding 1 to the highest seen result ID.
              start += 1
              response = self.call_unary_with_deadline(
                  rpc='GetServers', req=_GetServersRequest(start_server_id=start))
              for server in response.server:
                  start = max(start, server.ref.server_id)
                  yield server
  
      def list_server_sockets(self, server: Server) -> Iterator[Socket]:
          """List all server sockets that exist in server process.
  
          Iterating over the results will resolve additional pages automatically.
          """
          start: int = -1
          response: Optional[_GetServerSocketsResponse] = None
          while start < 0 or not response.end:
              # From proto: To request subsequent pages, the client generates this
              # value by adding 1 to the highest seen result ID.
              start += 1
              response = self.call_unary_with_deadline(
                  rpc='GetServerSockets',
                  req=_GetServerSocketsRequest(server_id=server.ref.server_id,
                                               start_socket_id=start))
              socket_ref: SocketRef
              for socket_ref in response.socket_ref:
                  start = max(start, socket_ref.socket_id)
                  # Yield actual socket
                  yield self.get_socket(socket_ref.socket_id)
  
      def list_channel_sockets(self, channel: Channel) -> Iterator[Socket]:
          """List all sockets of all subchannels of a given channel."""
          for subchannel in self.list_channel_subchannels(channel):
              yield from self.list_subchannels_sockets(subchannel)
  
      def list_channel_subchannels(self,
                                   channel: Channel) -> Iterator[Subchannel]:
          """List all subchannels of a given channel."""
          for subchannel_ref in channel.subchannel_ref:
              yield self.get_subchannel(subchannel_ref.subchannel_id)
  
      def list_subchannels_sockets(self,
                                   subchannel: Subchannel) -> Iterator[Socket]:
          """List all sockets of a given subchannel."""
          for socket_ref in subchannel.socket_ref:
              yield self.get_socket(socket_ref.socket_id)
  
      def get_subchannel(self, subchannel_id) -> Subchannel:
          """Return a single Subchannel, otherwise raises RpcError."""
          response: _GetSubchannelResponse = self.call_unary_with_deadline(
              rpc='GetSubchannel',
              req=_GetSubchannelRequest(subchannel_id=subchannel_id))
          return response.subchannel
  
      def get_socket(self, socket_id) -> Socket:
          """Return a single Socket, otherwise raises RpcError."""
          response: _GetSocketResponse = self.call_unary_with_deadline(
              rpc='GetSocket', req=_GetSocketRequest(socket_id=socket_id))
          return response.socket
  