content: package org.allaymc.server.pack;
  
  import com.google.gson.Gson;
  import com.google.gson.GsonBuilder;
  import com.google.gson.Strictness;
  import lombok.SneakyThrows;
  import lombok.extern.slf4j.Slf4j;
  import org.allaymc.api.pack.PackManifest;
  import org.allaymc.server.utils.JSONUtils;
  import org.apache.commons.lang3.RandomStringUtils;
  
  import javax.crypto.Cipher;
  import javax.crypto.spec.IvParameterSpec;
  import javax.crypto.spec.SecretKeySpec;
  import java.io.ByteArrayOutputStream;
  import java.io.FileOutputStream;
  import java.io.InputStreamReader;
  import java.nio.charset.StandardCharsets;
  import java.nio.file.Files;
  import java.nio.file.Path;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.zip.ZipEntry;
  import java.util.zip.ZipFile;
  import java.util.zip.ZipOutputStream;
  
  /**
   * A util for encrypting pack.
   *
   * @author daoge_cmd
   */
  @Slf4j
  public final class PackEncryptor {
      private static final Gson GSON = new GsonBuilder()
              .disableHtmlEscaping()
              .serializeNulls() // NOTICE: Null should be serialized!
              .setStrictness(Strictness.LENIENT)
              .create();
      private static final int KEY_LENGTH = 32;
      private static final byte[] VERSION = new byte[]{(byte) 0x00, (byte) 0x00, (byte) 0x00, (byte) 0x00};
      private static final byte[] MAGIC = new byte[]{(byte) 0xFC, (byte) 0xB9, (byte) 0xCF, (byte) 0x9B};
      private static final List<String> EXCLUDED_FILES = List.of(PackManifest.FILE_NAME, "pack_icon.png", "bug_pack_icon.png");
  
      public static String encrypt(Path inputPath, Path outputPath) {
          var key = RandomStringUtils.secure().nextAlphabetic(KEY_LENGTH);
          try (var inputZip = new ZipFile(inputPath.toString())) {
              encrypt0(inputZip, outputPath, key);
          } catch (Exception e) {
              log.error("Failed to encrypt pack", e);
          }
          return key;
      }
  
      public static void encrypt(Path inputPath, Path outputPath, String key) {
          try (var inputZip = new ZipFile(inputPath.toString())) {
              encrypt0(inputZip, outputPath, key);
          } catch (Exception e) {
              log.error("Failed to encrypt pack", e);
          }
      }
  
      @SneakyThrows
      private static void encrypt0(ZipFile inputZip, Path outputPath, String key) {
          if (key.length() != KEY_LENGTH) {
              throw new IllegalArgumentException("key length must be 32");
          }
  
          // Find content id
          var uuid = findPackUUID(inputZip);
  
          var contentEntries = new ArrayList<ContentEntry>();
  
          // Delete old output
          Files.deleteIfExists(outputPath);
          var outputStream = new ZipOutputStream(new FileOutputStream(outputPath.toFile()), StandardCharsets.UTF_8);
          // Encrypt files
          inputZip.stream().forEach(zipEntry -> {
              if (zipEntry.isDirectory()) {
                  createDirectoryRoot(zipEntry, outputStream);
                  if (isSubPackRoot(zipEntry)) {
                      // Handle sub pack
                      encryptSubPack(inputZip, outputStream, zipEntry.getName(), key, uuid);
                  }
  
                  return;
              }
              // Sub pack files will be handled in encryptSubPack()
              if (isSubPackFile(zipEntry)) {
                  return;
              }
  
              String entryKey = null;
              // Check if file is excluded
              if (EXCLUDED_FILES.contains(zipEntry.getName())) {
                  encryptExcludedFile(inputZip, outputStream, zipEntry);
                  // Excluded file does not have entry key
              } else {
                  // Encrypt file
                  entryKey = encryptFile(inputZip, outputStream, zipEntry);
              }
              contentEntries.add(new ContentEntry(zipEntry.getName(), entryKey));
          });
  
          generateContentsJson("contents.json", outputStream, uuid, key, contentEntries);
          outputStream.close();
      }
  
      @SneakyThrows
      private static void encryptSubPack(ZipFile inputZip, ZipOutputStream zos, String subPackPath, String key, String contentId) {
          var subPackContentEntries = new ArrayList<ContentEntry>();
  
          // Encrypt files
          inputZip.stream().forEach(zipEntry -> {
              if (zipEntry.isDirectory() || !zipEntry.getName().startsWith(subPackPath)) {
                  return;
              }
  
              String entryKey = encryptFile(inputZip, zos, zipEntry);
              subPackContentEntries.add(new ContentEntry(zipEntry.getName().substring(subPackPath.length()), entryKey));
          });
  
          generateContentsJson(subPackPath + "contents.json", zos, contentId, key, subPackContentEntries);
      }
  
      @SneakyThrows
      private static void generateContentsJson(String name, ZipOutputStream outputStream, String contentId, String key, ArrayList<ContentEntry> contentEntries) {
          outputStream.putNextEntry(new ZipEntry(name));
          try (var stream = new ByteArrayOutputStream()) {
              stream.write(VERSION);
              stream.write(MAGIC);
              paddingTo(stream, 0x10);
              var contentIdBytes = contentId.getBytes(StandardCharsets.UTF_8);
              // Write content id length
              stream.write(contentIdBytes.length);
              // Write content id
              stream.write(contentIdBytes);
              // Init contents.json encryptor
              var secretKey = new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), "AES");
              var cipher = Cipher.getInstance("AES/CFB8/NoPadding");
              cipher.init(Cipher.ENCRYPT_MODE, secretKey, new IvParameterSpec(key.substring(0, 16).getBytes(StandardCharsets.UTF_8)));
              // Write contents.json
              var contentJson = GSON.toJson(new Content(contentEntries));
              paddingTo(stream, 0x100);
              stream.write(cipher.doFinal(contentJson.getBytes(StandardCharsets.UTF_8)));
              outputStream.write(stream.toByteArray());
          }
          outputStream.closeEntry();
      }
  
      @SneakyThrows
      private static void encryptExcludedFile(ZipFile inputZip, ZipOutputStream outputStream, ZipEntry zipEntry) {
          outputStream.putNextEntry(copyZipEntry(zipEntry));
          outputStream.write(inputZip.getInputStream(zipEntry).readAllBytes());
          outputStream.closeEntry();
      }
  
      @SneakyThrows
      private static String encryptFile(ZipFile inputZip, ZipOutputStream outputStream, ZipEntry zipEntry) {
          byte[] bytes;
          bytes = inputZip.getInputStream(zipEntry).readAllBytes();
          // Init encryptor
          var key = RandomStringUtils.secure().nextAlphabetic(KEY_LENGTH);
          var secretKey = new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), "AES");
          var cipher = Cipher.getInstance("AES/CFB8/NoPadding");
          cipher.init(Cipher.ENCRYPT_MODE, secretKey, new IvParameterSpec(key.substring(0, 16).getBytes(StandardCharsets.UTF_8)));
          // Encrypt the file
          var encryptedBytes = cipher.doFinal(bytes);
          // Write bytes
          outputStream.putNextEntry(copyZipEntry(zipEntry));
          outputStream.write(encryptedBytes);
          outputStream.closeEntry();
          return key;
      }
  
      private static boolean isSubPackFile(ZipEntry zipEntry) {
          return zipEntry.getName().startsWith("subpacks/");
      }
  
      private static boolean isSubPackRoot(ZipEntry zipEntry) {
          return zipEntry.isDirectory() && zipEntry.getName().startsWith("subpacks/") && calculateCharCount(zipEntry.getName(), '/') == 2;
      }
  
      private static int calculateCharCount(String str, char target) {
          int count = 0;
          for (char c : str.toCharArray()) {
              if (c == target) {
                  count++;
              }
          }
          return count;
      }
  
      private static void paddingTo(ByteArrayOutputStream stream, int pos) {
          if (pos <= stream.size()) {
              throw new IllegalArgumentException("pos must be bigger than stream size");
          }
  
          var need = pos - stream.size();
          for (int i = 0; i < need; i++) {
              stream.write(0);
          }
      }
  
      private static ZipEntry copyZipEntry(ZipEntry entry) {
          var newEntry = new ZipEntry(entry);
          // Explicitly set method to DEFLATED to avoid invalid crc-32 error
          newEntry.setMethod(ZipEntry.DEFLATED);
          return newEntry;
      }
  
      @SneakyThrows
      private static String findPackUUID(ZipFile zip) {
          var manifestEntry = zip.getEntry(PackManifest.FILE_NAME);
          if (manifestEntry == null) {
              throw new IllegalArgumentException("manifest file not exists");
          }
  
          var manifest = JSONUtils.from(new InputStreamReader(zip.getInputStream(manifestEntry), StandardCharsets.UTF_8), PackManifest.class);
          return manifest.getHeader().getUuid().toString();
      }
  
      @SneakyThrows
      private static void createDirectoryRoot(ZipEntry zipEntry, ZipOutputStream outputStream) {
          outputStream.putNextEntry(copyZipEntry(zipEntry));
          outputStream.closeEntry();
      }
  
      public record Content(List<ContentEntry> content) {
      }
  
      public record ContentEntry(String path, String key) {
      }
  }
  