Add injection ability, boost tests, version bump 0.1.6

build.gradle

@@ -2,7 +2,7 @@ apply plugin: 'java'
 apply plugin: 'jacoco'
 apply plugin: 'com.github.kt3k.coveralls'
 
-version = '0.1.5'
+version = '0.1.6'
 
 repositories {
     mavenCentral()

src/main/java/android/content/res/AXMLResource.java

@@ -20,7 +20,9 @@
 import android.content.res.chunk.sections.ResourceSection;
 import android.content.res.chunk.sections.StringSection;
 import android.content.res.chunk.types.AXMLHeader;
+import android.content.res.chunk.types.Attribute;
 import android.content.res.chunk.types.Chunk;
+import android.content.res.chunk.types.StartTag;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -53,6 +55,29 @@ public AXMLResource(InputStream stream) throws IOException {
         }
     }
 
+    public void injectApplicationAttribute(Attribute attribute) {
+        StartTag tag = getApplicationTag();
+
+        tag.insertOrReplaceAttribute(attribute);
+    }
+
+    public StartTag getApplicationTag() {
+        Iterator<Chunk> iterator = chunks.iterator();
+        while (iterator.hasNext()) {
+            Chunk chunk = iterator.next();
+            if (chunk instanceof StartTag &&
+                    ((StartTag) chunk).getName(stringSection).equalsIgnoreCase("application")) {
+                return (StartTag) chunk;
+            }
+        }
+
+        return null;
+    }
+
+    public StringSection getStringSection() {
+        return stringSection;
+    }
+
     public boolean read(InputStream stream) throws IOException {
 
         IntReader reader = new IntReader(stream, false);

src/main/java/android/content/res/chunk/sections/GenericChunkSection.java

@@ -38,6 +38,7 @@ public GenericChunkSection(ChunkType chunkType, IntReader reader) {
 
             reader.skip(Math.abs(reader.getBytesRead() - getStartPosition() - size));
         } catch (IOException e) {
+            // Catching this here allows us to continue reading
             e.printStackTrace();
         }
     }

src/main/java/android/content/res/chunk/sections/StringSection.java

@@ -114,8 +114,31 @@ private void readPool(ArrayList<PoolItem> pool, int flags, IntReader inputReader
         }
     }
 
+    public int getStringIndex(String string) {
+        if (string != null) {
+            for (PoolItem item : stringChunkPool) {
+                if (item.getString().equals(string)) {
+                    return stringChunkPool.indexOf(item);
+                }
+            }
+        }
+
+        return -1;
+    }
+
+    public int putStringIndex(String string) {
+        int currentPosition = getStringIndex(string);
+        if (currentPosition != -1) {
+            return currentPosition;
+        }
+
+        stringChunkPool.add(new PoolItem(-1, string));
+
+        return getStringIndex(string);
+    }
+
     public String getString(int index) {
-        if ((index > -1) && (index < stringChunkCount)) {
+        if ((index > -1) && (index < stringChunkPool.size())) {
             return stringChunkPool.get(index).getString();
         }
 
@@ -151,14 +174,14 @@ public int getSize() {
 
         int styleDataSize = 0;
         for (PoolItem item : styleChunkPool) {
-            stringDataSize += item.getString().length() * (((stringChunkFlags & UTF8_FLAG) == 0) ? 2 : 1);
+            styleDataSize += item.getString().length() * (((stringChunkFlags & UTF8_FLAG) == 0) ? 2 : 1);
         }
 
         return (2 * 4) + // Header
                 (5 * 4) + // static sections
-                (stringChunkCount * 4) + // string table offset size
+                (stringChunkPool.size() * 4) + // string table offset size
                 stringDataSize +
-                (styleChunkCount * 4) + // style table offset size
+                (styleChunkPool.size() * 4) + // style table offset size
                 styleDataSize;
     }
 
@@ -247,21 +270,21 @@ public byte[] toBytes() {
         int newStringChunkOffset = 0;
         if (!stringChunkPool.isEmpty()) {
             newStringChunkOffset = (5 * 4) /* header + 3 other ints above it */
-                    + stringChunkCount * 4 /* index table size */
+                    + stringChunkPool.size() * 4 /* index table size */
                     + 8 /* (this space and the style chunk offset */;
         }
 
         int newStyleChunkOffset = 0;
         if (!styleChunkPool.isEmpty()) {
             newStyleChunkOffset = (6 * 4) /* header + 4 other ints above it */
-                    + styleChunkCount * 4 /* index table size */
+                    + styleChunkPool.size() * 4 /* index table size */
                     + 8 /* (this space and the style chunk offset */;
         }
 
         byte[] body = ByteBuffer.allocate(5 * 4)
                 .order(ByteOrder.LITTLE_ENDIAN)
-                .putInt(stringChunkCount)
-                .putInt(styleChunkCount)
+                .putInt(stringChunkPool.size())
+                .putInt(styleChunkPool.size())
                 .putInt(stringChunkFlags)
                 .putInt(newStringChunkOffset)
                 .putInt(newStyleChunkOffset)

src/main/java/android/content/res/chunk/types/Attribute.java

@@ -39,6 +39,28 @@ public class Attribute implements Chunk {
     private int attributeType;
     private int data;
 
+    public Attribute(String uri,
+                     String name,
+                     String stringData,
+                     AttributeType type,
+                     Object data,
+                     StringSection stringSection) {
+        this.uri = stringSection.getStringIndex(uri);
+        this.name = stringSection.getStringIndex(name);
+        this.stringData = stringSection.getStringIndex(stringData);
+        this.attributeType = type.getIntType();
+
+        if (attributeType == AttributeType.STRING.getIntType()) {
+            if (this.stringData == -1) {
+                this.stringData = stringSection.putStringIndex(stringData);
+            }
+            this.data = -1;
+        } else {
+            this.data = (int) data;
+        }
+
+    }
+
     public Attribute(IntReader reader) {
         try {
             uri = reader.readInt();
@@ -82,6 +104,14 @@ public int getSize() {
         return 4 * 5;
     }
 
+    public int getNameIndex() {
+        return name;
+    }
+
+    public int getStringDataIndex() {
+        return stringData;
+    }
+
     /*
      * (non-Javadoc)
      * 

src/main/java/android/content/res/chunk/types/StartTag.java

@@ -24,6 +24,7 @@
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.util.ArrayList;
+import java.util.Iterator;
 
 /**
  * StartTag type of Chunk, differs from a Namespace as there will be specific metadata inside of it
@@ -59,8 +60,9 @@ public void readHeader(IntReader inputReader) throws IOException {
         flags = inputReader.readInt();
         attributeCount = inputReader.readInt();
         classAttribute = inputReader.readInt();
+
+        attributes = new ArrayList<>();
         if (attributeCount > 0) {
-            attributes = new ArrayList<>();
             for (int i = 0; i < attributeCount; i++) {
                 attributes.add(new Attribute(inputReader));
             }
@@ -76,6 +78,26 @@ public int getSize() {
         return (9 * 4) + (attributeCount * 20);
     }
 
+    public ArrayList<Attribute> getAttributes() {
+        return attributes;
+    }
+
+    public void insertOrReplaceAttribute(Attribute newAttribute) {
+        Iterator<Attribute> iterator = attributes.iterator();
+        while (iterator.hasNext()) {
+            Attribute attribute = iterator.next();
+            if (attribute.getNameIndex() == newAttribute.getNameIndex()) {
+                iterator.remove();
+            }
+        }
+
+        attributes.add(newAttribute);
+    }
+
+    public String getName(StringSection stringSection) {
+        return stringSection.getString(name);
+    }
+
     /*
      * (non-Javadoc)
      * 
@@ -118,13 +140,13 @@ public byte[] toBytes() {
                 .putInt(namespaceUri)
                 .putInt(name)
                 .putInt(flags)
-                .putInt(attributeCount)
+                .putInt(attributes.size())
                 .putInt(classAttribute)
                 .array();
 
         byte[] dynamicBody;
-        if (attributeCount > 0) {
-            ByteBuffer attributeData = ByteBuffer.allocate(attributeCount * 20)
+        if (attributes.size() > 0) {
+            ByteBuffer attributeData = ByteBuffer.allocate(attributes.size() * 20)
                     .order(ByteOrder.LITTLE_ENDIAN);
             for (Attribute attribute : attributes) {
                 attributeData.put(attribute.toBytes());

src/test/java/android/content/res/TestAXMLResource.java

@@ -0,0 +1,104 @@
+package android.content.res;
+
+import android.content.res.chunk.AttributeType;
+import android.content.res.chunk.types.Attribute;
+import android.content.res.chunk.types.StartTag;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.experimental.runners.Enclosed;
+import org.junit.runner.RunWith;
+
+import java.io.*;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * @author tstrazzere
+ */
+@RunWith(Enclosed.class)
+public class TestAXMLResource {
+    public static class FunctionalTest {
+
+        // Legacy files from original repo
+        String[] oldTestFiles = {"test.xml", "test1.xml", "test2.xml", "test3.xml"};
+
+        // Large file with weird tricks that broke tools in the past
+        String largeFromMalware = "large_from_malware.xml";
+
+        AXMLResource underTest;
+
+        @Before
+        public void setUp() {
+            underTest = new AXMLResource();
+        }
+
+        @Test
+        public void testReadingOldFiles() throws IOException {
+            for (String file : oldTestFiles) {
+                InputStream testStream = this.getClass().getClassLoader().getResourceAsStream(file);
+
+                // Should throw no exceptions
+                underTest.read(testStream);
+            }
+        }
+
+        @Test
+        public void testPrinting() throws IOException {
+            InputStream testStream = this.getClass().getClassLoader().getResourceAsStream(largeFromMalware);
+
+            underTest = new AXMLResource(testStream);
+
+            underTest.print();
+        }
+
+        @Test
+        public void testInsertApplicationAttribute() throws IOException {
+            InputStream testStream = this.getClass().getClassLoader().getResourceAsStream(largeFromMalware);
+
+            underTest.read(testStream);
+
+            Attribute attribute = new Attribute("android",
+                    "name",
+                    "test",
+                    AttributeType.STRING,
+                    null,
+                    underTest.getStringSection());
+
+            underTest.injectApplicationAttribute(attribute);
+
+            StartTag startTag = underTest.getApplicationTag();
+
+            assertTrue(startTag.getAttributes().contains(attribute));
+        }
+
+        @Test
+        public void testWriteInsertedApplicationAttribute() throws IOException {
+            InputStream testStream = this.getClass().getClassLoader().getResourceAsStream(largeFromMalware);
+
+            underTest.read(testStream);
+
+            Attribute attribute = new Attribute("android",
+                    "name",
+                    "test",
+                    AttributeType.STRING,
+                    null,
+                    underTest.getStringSection());
+
+            underTest.injectApplicationAttribute(attribute);
+
+            File file = File.createTempFile("axml-func-test", "xml-test");
+            file.deleteOnExit();
+
+            underTest.write(new FileOutputStream(file));
+
+            underTest = new AXMLResource(new FileInputStream(file));
+            StartTag startTag = underTest.getApplicationTag();
+
+            assertEquals(underTest.getStringSection().getString(startTag.getAttributes().get(3).getNameIndex()),
+                    "name");
+            assertEquals(underTest.getStringSection().getString(startTag.getAttributes().get(3).getStringDataIndex()),
+                    "test");
+        }
+    }
+}