content: // Licensed to the Apache Software Foundation (ASF) under one
  // or more contributor license agreements.  See the NOTICE file
  // distributed with this work for additional information
  // regarding copyright ownership.  The ASF 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
  //
  //   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 org.apache.cloudstack.userdata;
  
  import java.io.ByteArrayInputStream;
  import java.io.ByteArrayOutputStream;
  import java.io.IOException;
  import java.io.InputStream;
  import java.util.Arrays;
  import java.util.List;
  import java.util.Map;
  import java.util.Properties;
  import java.util.stream.Collectors;
  import java.util.zip.GZIPInputStream;
  
  import javax.mail.BodyPart;
  import javax.mail.MessagingException;
  import javax.mail.Multipart;
  import javax.mail.Session;
  import javax.mail.internet.MimeBodyPart;
  import javax.mail.internet.MimeMessage;
  import javax.mail.internet.MimeMultipart;
  
  import org.apache.commons.codec.binary.Base64;
  import org.apache.commons.collections.CollectionUtils;
  import org.apache.commons.lang3.StringUtils;
  
  import com.cloud.utils.component.AdapterBase;
  import com.cloud.utils.exception.CloudRuntimeException;
  import com.sun.mail.util.BASE64DecoderStream;
  
  public class CloudInitUserDataProvider extends AdapterBase implements UserDataProvider {
  
      protected enum FormatType {
          CLOUD_CONFIG, BASH_SCRIPT, MIME, CLOUD_BOOTHOOK, INCLUDE_FILE
      }
  
      private static final String CLOUD_CONFIG_CONTENT_TYPE = "text/cloud-config";
      private static final String BASH_SCRIPT_CONTENT_TYPE = "text/x-shellscript";
      private static final String INCLUDE_FILE_CONTENT_TYPE = "text/x-include-url";
      private static final String CLOUD_BOOTHOOK_CONTENT_TYPE = "text/cloud-boothook";
  
      private static final Map<FormatType, String> formatContentTypeMap = Map.ofEntries(
              Map.entry(FormatType.CLOUD_CONFIG, CLOUD_CONFIG_CONTENT_TYPE),
              Map.entry(FormatType.BASH_SCRIPT, BASH_SCRIPT_CONTENT_TYPE),
              Map.entry(FormatType.CLOUD_BOOTHOOK, CLOUD_BOOTHOOK_CONTENT_TYPE),
              Map.entry(FormatType.INCLUDE_FILE, INCLUDE_FILE_CONTENT_TYPE)
      );
  
      private static final Session session = Session.getDefaultInstance(new Properties());
  
      @Override
      public String getName() {
          return "cloud-init";
      }
  
      protected boolean isGZipped(String encodedUserdata) {
          if (StringUtils.isEmpty(encodedUserdata)) {
              return false;
          }
          byte[] data = Base64.decodeBase64(encodedUserdata);
          if (data.length < 2) {
              return false;
          }
          int magic = data[0] & 0xff | ((data[1] << 8) & 0xff00);
          return magic == GZIPInputStream.GZIP_MAGIC;
      }
  
      protected String extractUserDataHeader(String userdata) {
          List<String> lines = Arrays.stream(userdata.split("\n"))
                  .filter(x -> (x.startsWith("#") && !x.startsWith("##")) || (x.startsWith("Content-Type:")))
                  .collect(Collectors.toList());
          if (CollectionUtils.isEmpty(lines)) {
              logger.debug("Failed to detect the user data format type as it does not contain a header");
              return null;
          }
          return lines.get(0);
      }
  
      protected FormatType mapUserDataHeaderToFormatType(String header, FormatType defaultFormatType) {
          if (StringUtils.isBlank(header)) {
              if (defaultFormatType == null) {
                  throw new CloudRuntimeException("Failed to detect the user data format type as it does not contain a header");
              }
              logger.debug(String.format("Empty header for userdata, using the default format type: %s", defaultFormatType.name()));
              return defaultFormatType;
          } else if (header.equalsIgnoreCase("#cloud-config")) {
              return FormatType.CLOUD_CONFIG;
          } else if (header.startsWith("#!")) {
              return FormatType.BASH_SCRIPT;
          } else if (header.equalsIgnoreCase("#cloud-boothook")) {
              return FormatType.CLOUD_BOOTHOOK;
          } else if (header.startsWith("#include")) {
              return FormatType.INCLUDE_FILE;
          } else if (header.startsWith("Content-Type:")) {
              return FormatType.MIME;
          } else {
              String msg = String.format("Cannot recognise the user data format type from the header line: %s." +
                      "Supported types are: cloud-config, bash script, cloud-boothook, include file or MIME", header);
              logger.error(msg);
              throw new CloudRuntimeException(msg);
          }
      }
  
      /**
       * Detect the user data type
       * @param userdata the userdata string to detect the type
       * @param defaultFormatType if not null, then use it in case the header does not exist in the userdata, otherwise fail
       * Reference: <a href="https://canonical-cloud-init.readthedocs-hosted.com/en/latest/explanation/format.html#user-data-formats" />
       */
      protected FormatType getUserDataFormatType(String userdata, FormatType defaultFormatType) {
          if (StringUtils.isBlank(userdata)) {
              String msg = "User data expected but provided empty user data";
              logger.error(msg);
              throw new CloudRuntimeException(msg);
          }
  
          String header = extractUserDataHeader(userdata);
          return mapUserDataHeaderToFormatType(header, defaultFormatType);
      }
  
      private String getContentType(String userData, FormatType formatType) throws MessagingException {
          if (formatType == FormatType.MIME) {
              NoIdMimeMessage msg = new NoIdMimeMessage(session, new ByteArrayInputStream(userData.getBytes()));
              return msg.getContentType();
          }
          if (!formatContentTypeMap.containsKey(formatType)) {
              throw new CloudRuntimeException(String.format("Cannot get the user data content type as " +
                      "its format type %s is invalid", formatType.name()));
          }
          return formatContentTypeMap.get(formatType);
      }
  
      protected String getBodyPartContentAsString(BodyPart bodyPart) throws MessagingException, IOException {
          Object content = bodyPart.getContent();
          if (content instanceof BASE64DecoderStream) {
              return new String(((BASE64DecoderStream)bodyPart.getContent()).readAllBytes());
          } else if (content instanceof ByteArrayInputStream) {
              return new String(((ByteArrayInputStream)bodyPart.getContent()).readAllBytes());
          } else if (content instanceof String) {
              return (String)bodyPart.getContent();
          }
          throw new CloudRuntimeException(String.format("Failed to get content for multipart data with content type: %s", getBodyPartContentType(bodyPart)));
      }
  
      private String getBodyPartContentType(BodyPart bodyPart) throws MessagingException {
          String contentType = StringUtils.defaultString(bodyPart.getDataHandler().getContentType(), bodyPart.getContentType());
          return  contentType.contains(";") ? contentType.substring(0, contentType.indexOf(';')) : contentType;
      }
  
      protected MimeBodyPart generateBodyPartMimeMessage(String userData, String contentType) throws MessagingException {
          MimeBodyPart bodyPart = new MimeBodyPart();
          bodyPart.setContent(userData, contentType);
          bodyPart.addHeader("Content-Transfer-Encoding", "base64");
          return bodyPart;
      }
  
      protected MimeBodyPart generateBodyPartMimeMessage(String userData, FormatType formatType) throws MessagingException {
          return generateBodyPartMimeMessage(userData, getContentType(userData, formatType));
      }
  
      private Multipart getMessageContent(NoIdMimeMessage message) {
          Multipart messageContent;
          try {
              messageContent = (MimeMultipart) message.getContent();
          } catch (IOException | MessagingException e) {
              messageContent = new MimeMultipart();
          }
          return messageContent;
      }
  
      private void addBodyPartToMultipart(Multipart existingMultipart, MimeBodyPart bodyPart) throws MessagingException, IOException {
          boolean added = false;
          final int existingCount = existingMultipart.getCount();
          for (int j = 0; j < existingCount; ++j) {
              MimeBodyPart existingBodyPart = (MimeBodyPart)existingMultipart.getBodyPart(j);
              String existingContentType = getBodyPartContentType(existingBodyPart);
              String newContentType = getBodyPartContentType(bodyPart);
              if (existingContentType.equals(newContentType)) {
                  String existingContent = getBodyPartContentAsString(existingBodyPart);
                  String newContent = getBodyPartContentAsString(bodyPart);
                  // generating a combined content MimeBodyPart to replace
                  MimeBodyPart combinedBodyPart = generateBodyPartMimeMessage(
                          simpleAppendSameFormatTypeUserData(existingContent, newContent), existingContentType);
                  existingMultipart.removeBodyPart(j);
                  existingMultipart.addBodyPart(combinedBodyPart, j);
                  added = true;
                  break;
              }
          }
          if (!added) {
              existingMultipart.addBodyPart(bodyPart);
          }
      }
  
      private void addBodyPartsToMessageContentFromUserDataContent(Multipart existingMultipart,
                                                                   NoIdMimeMessage msgFromUserdata) throws MessagingException, IOException {
          MimeMultipart newMultipart = (MimeMultipart)msgFromUserdata.getContent();
          final int existingCount = existingMultipart.getCount();
          final int newCount = newMultipart.getCount();
          for (int i = 0; i < newCount; ++i) {
              BodyPart bodyPart = newMultipart.getBodyPart(i);
              if (existingCount == 0) {
                  existingMultipart.addBodyPart(bodyPart);
                  continue;
              }
              addBodyPartToMultipart(existingMultipart, (MimeBodyPart)bodyPart);
          }
      }
  
      private NoIdMimeMessage createMultipartMessageAddingUserdata(String userData, FormatType formatType,
                                                             NoIdMimeMessage message) throws MessagingException, IOException {
          NoIdMimeMessage newMessage = new NoIdMimeMessage(session);
          Multipart messageContent = getMessageContent(message);
  
          if (formatType == FormatType.MIME) {
              NoIdMimeMessage msgFromUserdata = new NoIdMimeMessage(session, new ByteArrayInputStream(userData.getBytes()));
              addBodyPartsToMessageContentFromUserDataContent(messageContent, msgFromUserdata);
          } else {
              MimeBodyPart part = generateBodyPartMimeMessage(userData, formatType);
              addBodyPartToMultipart(messageContent, part);
          }
          newMessage.setContent(messageContent);
          return newMessage;
      }
  
      private String simpleAppendSameFormatTypeUserData(String userData1, String userData2) {
          String userdata2Header = extractUserDataHeader(userData2);
          int beginIndex = StringUtils.isNotBlank(userdata2Header) ? userData2.indexOf('\n')+1 : 0;
          return String.format("%s\n\n%s", userData1, userData2.substring(beginIndex));
      }
  
      private void checkGzipAppend(String encodedUserData1, String encodedUserData2) {
          if (isGZipped(encodedUserData1) || isGZipped(encodedUserData2)) {
              throw new CloudRuntimeException("Gzipped user data can not be used together with other user data formats");
          }
      }
  
      @Override
      public String appendUserData(String encodedUserData1, String encodedUserData2) {
          try {
              checkGzipAppend(encodedUserData1, encodedUserData2);
              String userData1 = new String(Base64.decodeBase64(encodedUserData1));
              String userData2 = new String(Base64.decodeBase64(encodedUserData2));
              FormatType formatType1 = getUserDataFormatType(userData1, null);
              FormatType formatType2 = getUserDataFormatType(userData2, formatType1);
              if (formatType1.equals(formatType2) && List.of(FormatType.CLOUD_CONFIG, FormatType.BASH_SCRIPT).contains(formatType1)) {
                  return simpleAppendSameFormatTypeUserData(userData1, userData2);
              }
              NoIdMimeMessage message = new NoIdMimeMessage(session);
              message = createMultipartMessageAddingUserdata(userData1, formatType1, message);
              message = createMultipartMessageAddingUserdata(userData2, formatType2, message);
              ByteArrayOutputStream output = new ByteArrayOutputStream();
              message.writeTo(output);
              return output.toString();
          } catch (MessagingException | IOException | CloudRuntimeException e) {
              String msg = String.format("Error attempting to merge user data as a multipart user data. " +
                      "Reason: %s", e.getMessage());
              logger.error(msg, e);
              throw new CloudRuntimeException(msg, e);
          }
      }
  
      /* This is a wrapper class just to remove Message-ID header from the resultant
         multipart data which may contain server details.
       */
      private class NoIdMimeMessage extends MimeMessage {
          NoIdMimeMessage (Session session) {
              super(session);
          }
          NoIdMimeMessage (Session session, InputStream is) throws MessagingException {
              super(session, is);
          }
          @Override
          protected void updateMessageID() throws MessagingException {
              removeHeader("Message-ID");
          }
      }
  }
  