/*
 * JBoss, Home of Professional Open Source.
 * Copyright 2014 Red Hat, Inc., and individual contributors
 * as indicated by the @author tags.
 *
 * 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.
 */

package io.netty.handler.ssl;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

import javax.net.ssl.SSLException;

/**
 * Hacks up ALPN support into the server hello message
 *
 * This has two different usage modes, one is adding a selected protocol into the extensions, the other is removing
 * all mention of ALPN and retuning the selected protocol. This dual mode does not make for the cleanest code
 * but removes the need to have duplicate nearly identical methods.
 *
 * The if the selected protocol is set then this will be added. If the selected protocol is null then ALPN will be
 * parsed and removed.
 *
 * <p>
 * We only care about TLS 1.2, as TLS 1.1 is not allowed to use ALPN.
 * <p>
 * Super hacky, but slightly less hacky than modifying the boot class path
 */
final class ALPNHackServerHelloExplorer {

   // Private constructor prevents construction outside this class.
   private ALPNHackServerHelloExplorer() {
   }

   static byte[] addAlpnExtensionsToServerHello(byte[] source, String selectedAlpnProtocol)
         throws SSLException {
      ByteArrayOutputStream out = new ByteArrayOutputStream();
      ByteBuffer input = ByteBuffer.wrap(source);
      try {

         exploreHandshake(input, source.length, new AtomicReference<>(selectedAlpnProtocol), out);
         //we need to adjust the record length;
         int serverHelloLength = out.size() - 4;
         out.write(source, input.position(), input.remaining()); //there may be more messages (cert etc), so we append them
         byte[] data = out.toByteArray();

         //now we need to adjust the handshake frame length
         data[1] = (byte) ((serverHelloLength >> 16) & 0xFF);
         data[2] = (byte) ((serverHelloLength >> 8) & 0xFF);
         data[3] = (byte) (serverHelloLength & 0xFF);
         return data;
      } catch (AlpnProcessingException e) {
         return source;
      }
   }

   /**
    * removes the ALPN extensions from the server hello
    * @param source
    * @return
    * @throws SSLException
    */
   static byte[] removeAlpnExtensionsFromServerHello(ByteBuffer source, final AtomicReference<String> selectedAlpnProtocol)
         throws SSLException {
      ByteArrayOutputStream out = new ByteArrayOutputStream();

      try {

         exploreHandshake(source, source.remaining(), selectedAlpnProtocol, out);
         //we need to adjust the record length;
         int serverHelloLength = out.size() - 4;
         byte[] data = out.toByteArray();

         //now we need to adjust the handshake frame length
         data[1] = (byte) ((serverHelloLength >> 16) & 0xFF);
         data[2] = (byte) ((serverHelloLength >> 8) & 0xFF);
         data[3] = (byte) (serverHelloLength & 0xFF);
         return data;
      } catch (AlpnProcessingException e) {
         return null;
      }
   }
   private static void exploreHandshake(ByteBuffer input, int recordLength, AtomicReference<String> selectedAlpnProtocol, ByteArrayOutputStream out) throws SSLException {

      // What is the handshake type?
      byte handshakeType = input.get();
      if (handshakeType != 0x02) {   // 0x01: server_hello message
         throw new SSLException("Expected server hello");
      }
      out.write(handshakeType);

      // What is the handshake body length?
      int handshakeLength = getInt24(input);
      out.write(0); //placeholders
      out.write(0);
      out.write(0);

      // Theoretically, a single handshake message might span multiple
      // records, but in practice this does not occur.
      if (handshakeLength > recordLength - 4) { // 4: handshake header size
         throw new SSLException("Multi record handshake");
      }
      int old = input.limit();
      input.limit(handshakeLength + input.position());
      exploreServerHello(input, selectedAlpnProtocol, out);
      input.limit(old);
   }

   private static void exploreServerHello( ByteBuffer input, AtomicReference<String> alpnProtocolReference, ByteArrayOutputStream out) throws SSLException {

      // server version
      byte helloMajorVersion = input.get();
      byte helloMinorVersion = input.get();
      out.write(helloMajorVersion);
      out.write(helloMinorVersion);

      for (int i = 0; i < 32; ++i) { //the Random is 32 bytes
         out.write(input.get() & 0xFF);
      }

      // ignore session id
      processByteVector8(input, out);

      // ignore cipher_suite
      out.write(input.get() & 0xFF);
      out.write(input.get() & 0xFF);

      // ignore compression methods
      out.write(input.get() & 0xFF);

      String existingAlpn = null;
      ByteArrayOutputStream extensionsOutput = null;
      if (input.remaining() > 0) {
         extensionsOutput = new ByteArrayOutputStream();
         existingAlpn = exploreExtensions(input, extensionsOutput, alpnProtocolReference.get() == null);
      }

      if (existingAlpn != null) {
         if(alpnProtocolReference.get() != null) {
            throw new AlpnProcessingException();
         }
         alpnProtocolReference.set(existingAlpn);
         byte[] existing = extensionsOutput.toByteArray();
         out.write(existing, 0, existing.length);

      } else if(alpnProtocolReference.get() != null) {
         String selectedAlpnProtocol = alpnProtocolReference.get();
         ByteArrayOutputStream alpnBits = new ByteArrayOutputStream();
         alpnBits.write(0);
         alpnBits.write(16); //ALPN type
         int length = 3 + selectedAlpnProtocol.length(); //length of extension data
         alpnBits.write((length >> 8) & 0xFF);
         alpnBits.write(length & 0xFF);
         length -= 2;
         alpnBits.write((length >> 8) & 0xFF);
         alpnBits.write(length & 0xFF);
         alpnBits.write(selectedAlpnProtocol.length() & 0xFF);
         for (int i = 0; i < selectedAlpnProtocol.length(); ++i) {
            alpnBits.write(selectedAlpnProtocol.charAt(i) & 0xFF);
         }

         if (extensionsOutput != null) {
            byte[] existing = extensionsOutput.toByteArray();
            int newLength = existing.length - 2 + alpnBits.size();
            existing[0] = (byte) ((newLength >> 8) & 0xFF);
            existing[1] = (byte) (newLength & 0xFF);
            try {
               out.write(existing);
               out.write(alpnBits.toByteArray());
            } catch (IOException e) {
               throw new RuntimeException(e);
            }
         } else {
            int al = alpnBits.size();
            out.write((al >> 8) & 0xFF);
            out.write(al & 0xFF);
            try {
               out.write(alpnBits.toByteArray());
            } catch (IOException e) {
               throw new RuntimeException(e);
            }
         }
      } else if(extensionsOutput != null){
         byte[] existing = extensionsOutput.toByteArray();
         out.write(existing, 0, existing.length);
      }
   }

   static List<ByteBuffer> extractRecords(ByteBuffer data) {
      List<ByteBuffer> ret = new ArrayList<>();
      while (data.hasRemaining()) {
         byte d1 = data.get();
         byte d2 = data.get();
         byte d3 = data.get();
         byte d4 = data.get();
         byte d5 = data.get();
         int length = (d4 & 0xFF) << 8 | d5 & 0xFF;
         byte[] b = new byte[length + 5];
         b[0] = d1;
         b[1] = d2;
         b[2] = d3;
         b[3] = d4;
         b[4] = d5;
         data.get(b, 5, length);
         ret.add(ByteBuffer.wrap(b));
      }
      return ret;
   }

   private static String exploreExtensions(ByteBuffer input, ByteArrayOutputStream extensionOut, boolean removeAlpn)
         throws SSLException {
      ByteArrayOutputStream out = new ByteArrayOutputStream();
      String ret = null;
      int length = getInt16(input);           // length of extensions
      out.write((length >> 8) & 0xFF);
      out.write(length & 0xFF);
      int originalLength = length;
      while (length > 0) {
         int extType = getInt16(input);      // extenson type
         int extLen = getInt16(input);       // length of extension data
         if(extType == 16) {
            int vlen = getInt16(input);
            ret = readByteVector8(input);
            if(!removeAlpn) {
               //we write the extension data back to the output stream
               out.write((extType >> 8) & 0xFF);
               out.write(extType & 0xFF);
               out.write((extLen >> 8) & 0xFF);
               out.write(extLen & 0xFF);
               out.write((vlen >> 8) & 0xFF);
               out.write(vlen & 0xFF);
               out.write(ret.length() & 0xFF);
               for(int i = 0; i < ret.length(); ++i) {
                  out.write(ret.charAt(i) & 0xFF);
               }
            } else {
               originalLength -= 6;
               originalLength -= vlen;
            }
         } else {
            out.write((extType >> 8) & 0xFF);
            out.write(extType & 0xFF);
            out.write((extLen >> 8) & 0xFF);
            out.write(extLen & 0xFF);
            processByteVector(input, extLen, out);
         }
         length -= extLen + 4;
      }
      if(removeAlpn && ret == null) {
         //there was not ALPN to remove, so this whole thing is unnecessary, throw an exception to abort
         throw new AlpnProcessingException();
      }
      byte[] data = out.toByteArray();
      data[0] = (byte) ((originalLength >> 8) & 0xFF);
      data[1] = (byte) (originalLength  & 0xFF);
      extensionOut.write(data, 0, data.length);
      return ret;
   }

   private static String readByteVector8(ByteBuffer input) {
      int length = getInt8(input);
      byte[] data = new byte[length];
      input.get(data);
      return new String(data, StandardCharsets.US_ASCII);
   }
   private static int getInt8(ByteBuffer input) {
      return input.get();
   }

   private static int getInt16(ByteBuffer input) {
      return (input.get() & 0xFF) << 8 | input.get() & 0xFF;
   }

   private static int getInt24(ByteBuffer input) {
      return (input.get() & 0xFF) << 16 | (input.get() & 0xFF) << 8 |
            input.get() & 0xFF;
   }
   private static void processByteVector8(ByteBuffer input, ByteArrayOutputStream out) {
      int int8 = getInt8(input);
      out.write(int8 & 0xFF);
      processByteVector(input, int8, out);
   }


   private static void processByteVector(ByteBuffer input, int length, ByteArrayOutputStream out) {
      for (int i = 0; i < length; ++i) {
         out.write(input.get() & 0xFF);
      }
   }

   static ByteBuffer createNewOutputRecords(byte[] newFirstMessage, List<ByteBuffer> records) {
      int length = newFirstMessage.length;
      length += 5; //Framing layer
      for (int i = 1; i < records.size(); ++i) {
         //the first record is the old server hello, so we start at 1 rather than zero
         ByteBuffer rec = records.get(i);
         length += rec.remaining();
      }
      byte[] newData = new byte[length];
      ByteBuffer ret = ByteBuffer.wrap(newData);
      ByteBuffer oldHello = records.get(0);
      ret.put(oldHello.get()); //type
      ret.put(oldHello.get()); //major
      ret.put(oldHello.get()); //minor
      ret.put((byte) ((newFirstMessage.length >> 8) & 0xFF));
      ret.put((byte) (newFirstMessage.length & 0xFF));
      ret.put(newFirstMessage);
      for (int i = 1; i < records.size(); ++i) {
         ByteBuffer rec = records.get(i);
         ret.put(rec);
      }
      ret.flip();
      return ret;
   }

   private static final class AlpnProcessingException extends RuntimeException {

   }
}
