/*
    * Copyright 2020 The Netty Project
    *
    * The Netty Project licenses this file to you 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:
    *
    *   https://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.
   */
  package io.netty.example.dns.tcp;
  
  import io.netty.bootstrap.Bootstrap;
  import io.netty.buffer.ByteBufUtil;
  
  import io.netty.channel.Channel;
  import io.netty.channel.ChannelInitializer;
  import io.netty.channel.ChannelHandlerContext;
  import io.netty.channel.ChannelPipeline;
  import io.netty.channel.EventLoopGroup;
  import io.netty.channel.nio.NioEventLoopGroup;
  import io.netty.channel.socket.SocketChannel;
  import io.netty.channel.socket.nio.NioSocketChannel;
  import io.netty.channel.SimpleChannelInboundHandler;
  
  import io.netty.handler.codec.dns.DefaultDnsQuestion;
  import io.netty.handler.codec.dns.DefaultDnsResponse;
  import io.netty.handler.codec.dns.DnsQuestion;
  import io.netty.handler.codec.dns.DnsQuery;
  import io.netty.handler.codec.dns.DefaultDnsQuery;
  import io.netty.handler.codec.dns.DnsOpCode;
  import io.netty.handler.codec.dns.DnsRecord;
  import io.netty.handler.codec.dns.DnsSection;
  import io.netty.handler.codec.dns.DnsRecordType;
  import io.netty.handler.codec.dns.DnsRawRecord;
  import io.netty.handler.codec.dns.TcpDnsQueryEncoder;
  import io.netty.handler.codec.dns.TcpDnsResponseDecoder;
  import io.netty.util.NetUtil;
  
  import java.util.Random;
  import java.util.concurrent.TimeUnit;
  
  public final class TcpDnsClient {
      private static final String QUERY_DOMAIN = "www.example.com";
      private static final int DNS_SERVER_PORT = 53;
      private static final String DNS_SERVER_HOST = "8.8.8.8";
  
      private TcpDnsClient() {
      }
  
      private static void handleQueryResp(DefaultDnsResponse msg) {
          if (msg.count(DnsSection.QUESTION) > 0) {
              DnsQuestion question = msg.recordAt(DnsSection.QUESTION, 0);
              System.out.printf("name: %s%n", question.name());
          }
          for (int i = 0, count = msg.count(DnsSection.ANSWER); i < count; i++) {
              DnsRecord record = msg.recordAt(DnsSection.ANSWER, i);
              if (record.type() == DnsRecordType.A) {
                  //just print the IP after query
                  DnsRawRecord raw = (DnsRawRecord) record;
                  System.out.println(NetUtil.bytesToIpAddress(ByteBufUtil.getBytes(raw.content())));
              }
          }
      }
  
      public static void main(String[] args) throws Exception {
          EventLoopGroup group = new NioEventLoopGroup();
          try {
              Bootstrap b = new Bootstrap();
              b.group(group)
                      .channel(NioSocketChannel.class)
                      .handler(new ChannelInitializer<SocketChannel>() {
                          @Override
                          protected void initChannel(SocketChannel ch) {
                              ChannelPipeline p = ch.pipeline();
                              p.addLast(new TcpDnsQueryEncoder())
                                      .addLast(new TcpDnsResponseDecoder())
                                      .addLast(new SimpleChannelInboundHandler<DefaultDnsResponse>() {
                                          @Override
                                          protected void channelRead0(ChannelHandlerContext ctx, DefaultDnsResponse msg) {
                                              try {
                                                  handleQueryResp(msg);
                                              } finally {
                                                  ctx.close();
                                              }
                                          }
                                      });
                          }
                      });
  
              final Channel ch = b.connect(DNS_SERVER_HOST, DNS_SERVER_PORT).sync().channel();
  
              int randomID = new Random().nextInt(60000 - 1000) + 1000;
              DnsQuery query = new DefaultDnsQuery(randomID, DnsOpCode.QUERY)
                     .setRecord(DnsSection.QUESTION, new DefaultDnsQuestion(QUERY_DOMAIN, DnsRecordType.A));
             ch.writeAndFlush(query).sync();
             boolean success = ch.closeFuture().await(10, TimeUnit.SECONDS);
             if (!success) {
                 System.err.println("dns query timeout!");
                 ch.close().sync();
             }
         } finally {
             group.shutdownGracefully();
         }
     }
 }