From abeb0b098466b70d8b5eaf1fb03339967f5027b6 Mon Sep 17 00:00:00 2001
From: cyian-1756 <devnull64@vfemail.net>
Date: Wed, 25 Jul 2018 03:16:41 -0400
Subject: [PATCH] Changed how magic number look ups are done (now uses a hash
 map); added magic number for png; added unit test

---
 .../java/com/rarchives/ripme/utils/Utils.java | 22 ++++++++++++++-----
 .../com/rarchives/ripme/tst/UtilsTest.java    |  3 +++
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/main/java/com/rarchives/ripme/utils/Utils.java b/src/main/java/com/rarchives/ripme/utils/Utils.java
index 31af93a0..1f49b8a5 100644
--- a/src/main/java/com/rarchives/ripme/utils/Utils.java
+++ b/src/main/java/com/rarchives/ripme/utils/Utils.java
@@ -20,6 +20,7 @@ import java.lang.reflect.Constructor;
 import java.net.URISyntaxException;
 import java.net.URL;
 import java.net.URLDecoder;
+import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Enumeration;
@@ -45,6 +46,7 @@ public class Utils {
 
     private static PropertiesConfiguration config;
     private static HashMap<String, HashMap<String, String>> cookieCache;
+    private static HashMap<ByteBuffer, String> magicHash = new HashMap<>();
 
     static {
         cookieCache = new HashMap<>();
@@ -732,13 +734,21 @@ public class Utils {
                 Utils.bytesToHumanReadable(bytesTotal);
     }
 
-    public static String getEXTFromMagic(byte[] magic) {
-        if (Arrays.equals(magic, new byte[]{-1, -40, -1, -37, 0, 0, 0, 0})) {
-            return "jpeg";
-        } else {
-            LOGGER.info("Unknown magic number " + Arrays.toString(magic));
+    public static String getEXTFromMagic(ByteBuffer magic) {
+        if (magicHash.isEmpty()) {
+            LOGGER.debug("initialising map");
+            initialiseMagicHashMap();
         }
-        return null;
+        return magicHash.get(magic);
+    }
+
+    public static String getEXTFromMagic(byte[] magic) {
+        return getEXTFromMagic(ByteBuffer.wrap(magic));
+    }
+
+    private static void initialiseMagicHashMap() {
+        magicHash.put(ByteBuffer.wrap(new byte[]{-1, -40, -1, -37, 0, 0, 0, 0}), "jpeg");
+        magicHash.put(ByteBuffer.wrap(new byte[]{-119, 80, 78, 71, 13, 0, 0, 0}), "png");
     }
 
     // Checks if a file exists ignoring it's extension.
diff --git a/src/test/java/com/rarchives/ripme/tst/UtilsTest.java b/src/test/java/com/rarchives/ripme/tst/UtilsTest.java
index d0789b39..66373274 100644
--- a/src/test/java/com/rarchives/ripme/tst/UtilsTest.java
+++ b/src/test/java/com/rarchives/ripme/tst/UtilsTest.java
@@ -3,10 +3,13 @@ package com.rarchives.ripme.tst;
 import junit.framework.TestCase;
 import com.rarchives.ripme.utils.Utils;
 
+import java.nio.ByteBuffer;
+
 public class UtilsTest extends TestCase {
 
     public void testGetEXTFromMagic() {
         assertEquals("jpeg", Utils.getEXTFromMagic(new byte[]{-1, -40, -1, -37, 0, 0, 0, 0}));
+        assertEquals("png", Utils.getEXTFromMagic(new byte[]{-119, 80, 78, 71, 13, 0, 0, 0}));
     }
 
     public void testStripURLParameter() {