[Forensic] Refactor ForensicEvent

ForensicEvent will contain events of type SecurityEvent,
ConnectEvent, or DnsEvent. The contained event will be
denoted by the string mType, and constrained to values
EventType.

Bug: 369313906
Test: m -j ; atest ForensicServiceTest
Flag: android.security.afl_api
Ignore-AOSP-First: security feature
Change-Id: Ic87aca38631ef09d5cd0a43b0cbe5f5ff204af5b
diff --git a/core/java/android/security/forensic/ForensicEvent.java b/core/java/android/security/forensic/ForensicEvent.java
index 90906ed..3d908cc 100644
--- a/core/java/android/security/forensic/ForensicEvent.java
+++ b/core/java/android/security/forensic/ForensicEvent.java
@@ -17,13 +17,17 @@
 package android.security.forensic;
 
 import android.annotation.FlaggedApi;
+import android.annotation.IntDef;
 import android.annotation.NonNull;
+import android.app.admin.ConnectEvent;
+import android.app.admin.DnsEvent;
+import android.app.admin.SecurityLog.SecurityEvent;
 import android.os.Parcel;
 import android.os.Parcelable;
 import android.security.Flags;
-import android.util.ArrayMap;
 
-import java.util.Map;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
 
 /**
  * A class that represents a forensic event.
@@ -33,11 +37,24 @@
 public final class ForensicEvent implements Parcelable {
     private static final String TAG = "ForensicEvent";
 
-    @NonNull
-    private final String mType;
+    public static final int SECURITY_EVENT = 0;
+    public static final int NETWORK_EVENT_DNS = 1;
+    public static final int NETWORK_EVENT_CONNECT = 2;
 
-    @NonNull
-    private final Map<String, String> mKeyValuePairs;
+    /** @hide */
+    @Retention(RetentionPolicy.SOURCE)
+    @IntDef({
+        ForensicEvent.SECURITY_EVENT,
+        ForensicEvent.NETWORK_EVENT_DNS,
+        ForensicEvent.NETWORK_EVENT_CONNECT,
+    })
+    public @interface EventType {}
+
+    @NonNull @EventType private final int mType;
+
+    private final SecurityEvent mSecurityEvent;
+    private final DnsEvent mNetworkEventDns;
+    private final ConnectEvent mNetworkEventConnect;
 
     public static final @NonNull Parcelable.Creator<ForensicEvent> CREATOR =
             new Parcelable.Creator<>() {
@@ -50,30 +67,99 @@
                 }
             };
 
-    public ForensicEvent(@NonNull String type, @NonNull Map<String, String> keyValuePairs) {
-        mType = type;
-        mKeyValuePairs = keyValuePairs;
+    public ForensicEvent(@NonNull SecurityEvent securityEvent) {
+        mType = SECURITY_EVENT;
+        mSecurityEvent = securityEvent;
+        mNetworkEventDns = null;
+        mNetworkEventConnect = null;
+    }
+
+    public ForensicEvent(@NonNull DnsEvent dnsEvent) {
+        mType = NETWORK_EVENT_DNS;
+        mNetworkEventDns = dnsEvent;
+        mSecurityEvent = null;
+        mNetworkEventConnect = null;
+    }
+
+    public ForensicEvent(@NonNull ConnectEvent connectEvent) {
+        mType = NETWORK_EVENT_CONNECT;
+        mNetworkEventConnect = connectEvent;
+        mSecurityEvent = null;
+        mNetworkEventDns = null;
     }
 
     private ForensicEvent(@NonNull Parcel in) {
-        mType = in.readString();
-        mKeyValuePairs = new ArrayMap<>(in.readInt());
-        in.readMap(mKeyValuePairs, getClass().getClassLoader(), String.class, String.class);
+        mType = in.readInt();
+        switch (mType) {
+            case SECURITY_EVENT:
+                mSecurityEvent = SecurityEvent.CREATOR.createFromParcel(in);
+                mNetworkEventDns = null;
+                mNetworkEventConnect = null;
+                break;
+            case NETWORK_EVENT_DNS:
+                mNetworkEventDns = DnsEvent.CREATOR.createFromParcel(in);
+                mSecurityEvent = null;
+                mNetworkEventConnect = null;
+                break;
+            case NETWORK_EVENT_CONNECT:
+                mNetworkEventConnect = ConnectEvent.CREATOR.createFromParcel(in);
+                mSecurityEvent = null;
+                mNetworkEventDns = null;
+                break;
+            default:
+                throw new IllegalArgumentException("Invalid event type: " + mType);
+        }
     }
 
-    public String getType() {
+    /** Returns the type of the forensic event. */
+    @NonNull
+    public @EventType int getType() {
         return mType;
     }
 
-    public Map<String, String> getKeyValuePairs() {
-        return mKeyValuePairs;
+    /** Returns the SecurityEvent object. */
+    @NonNull
+    public SecurityEvent getSecurityEvent() {
+        if (mType == SECURITY_EVENT) {
+            return mSecurityEvent;
+        }
+        throw new IllegalArgumentException("Event type is not security event: " + mType);
+    }
+
+    /** Returns the DnsEvent object. */
+    @NonNull
+    public DnsEvent getDnsEvent() {
+        if (mType == NETWORK_EVENT_DNS) {
+            return mNetworkEventDns;
+        }
+        throw new IllegalArgumentException("Event type is not network DNS event: " + mType);
+    }
+
+    /** Returns the ConnectEvent object. */
+    @NonNull
+    public ConnectEvent getConnectEvent() {
+        if (mType == NETWORK_EVENT_CONNECT) {
+            return mNetworkEventConnect;
+        }
+        throw new IllegalArgumentException("Event type is not network connect event: " + mType);
     }
 
     @Override
     public void writeToParcel(@NonNull Parcel out, int flags) {
-        out.writeString(mType);
-        out.writeInt(mKeyValuePairs.size());
-        out.writeMap(mKeyValuePairs);
+        out.writeInt(mType);
+        switch (mType) {
+            case SECURITY_EVENT:
+                out.writeParcelable(mSecurityEvent, flags);
+                break;
+            case NETWORK_EVENT_DNS:
+                out.writeParcelable(mNetworkEventDns, flags);
+                break;
+            case NETWORK_EVENT_CONNECT:
+                out.writeParcelable(mNetworkEventConnect, flags);
+                break;
+            default:
+                throw new IllegalArgumentException("Invalid event type: " + mType);
+        }
     }
 
     @FlaggedApi(Flags.FLAG_AFL_API)
@@ -86,7 +172,6 @@
     public String toString() {
         return "ForensicEvent{"
                 + "mType=" + mType
-                + ", mKeyValuePairs=" + mKeyValuePairs
                 + '}';
     }
 }
diff --git a/services/core/java/com/android/server/security/forensic/SecurityLogSource.java b/services/core/java/com/android/server/security/forensic/SecurityLogSource.java
index 0f1aa42..e1b49c4 100644
--- a/services/core/java/com/android/server/security/forensic/SecurityLogSource.java
+++ b/services/core/java/com/android/server/security/forensic/SecurityLogSource.java
@@ -22,9 +22,7 @@
 import android.app.admin.SecurityLog.SecurityEvent;
 import android.content.Context;
 import android.security.forensic.ForensicEvent;
-import android.util.ArrayMap;
 
-import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
@@ -34,10 +32,6 @@
 public class SecurityLogSource implements DataSource {
 
     private static final String TAG = "Forensic SecurityLogSource";
-    private static final String EVENT_TYPE = "SecurityEvent";
-    private static final String EVENT_TAG = "TAG";
-    private static final String EVENT_TIME = "TIME";
-    private static final String EVENT_DATA = "DATA";
 
     private SecurityEventCallback mEventCallback = new SecurityEventCallback();
     private DevicePolicyManager mDpm;
@@ -94,46 +88,9 @@
             List<ForensicEvent> forensicEvents =
                     events.stream()
                             .filter(event -> event != null)
-                            .map(event -> toForensicEvent(event))
+                            .map(event -> new ForensicEvent(event))
                             .collect(Collectors.toList());
             mDataAggregator.addBatchData(forensicEvents);
         }
-
-        private ForensicEvent toForensicEvent(SecurityEvent event) {
-            ArrayMap<String, String> keyValuePairs = new ArrayMap<>();
-            keyValuePairs.put(EVENT_TIME, String.valueOf(event.getTimeNanos()));
-            // TODO: Map tag to corresponding string
-            keyValuePairs.put(EVENT_TAG, String.valueOf(event.getTag()));
-            keyValuePairs.put(EVENT_DATA, eventDataToString(event.getData()));
-            return new ForensicEvent(EVENT_TYPE, keyValuePairs);
-        }
-
-        /**
-         * Convert event data to a String.
-         *
-         * @param obj Object containing an Integer, Long, Float, String, null, or Object[] of the
-         *     same.
-         * @return String representation of event data.
-         */
-        private String eventDataToString(Object obj) {
-            if (obj == null) {
-                return "";
-            } else if (obj instanceof Integer
-                    || obj instanceof Long
-                    || obj instanceof Float
-                    || obj instanceof String) {
-                return String.valueOf(obj);
-            } else if (obj instanceof Object[]) {
-                Object[] objArray = (Object[]) obj;
-                String[] strArray = new String[objArray.length];
-                for (int i = 0; i < objArray.length; ++i) {
-                    strArray[i] = eventDataToString(objArray[i]);
-                }
-                return Arrays.toString((String[]) strArray);
-            } else {
-                throw new IllegalArgumentException(
-                        "Unsupported data type: " + obj.getClass().getSimpleName());
-            }
-        }
     }
 }
diff --git a/services/tests/security/forensic/src/com/android/server/security/forensic/ForensicServiceTest.java b/services/tests/security/forensic/src/com/android/server/security/forensic/ForensicServiceTest.java
index 0da6db6..03c449c 100644
--- a/services/tests/security/forensic/src/com/android/server/security/forensic/ForensicServiceTest.java
+++ b/services/tests/security/forensic/src/com/android/server/security/forensic/ForensicServiceTest.java
@@ -31,6 +31,9 @@
 import static org.mockito.Mockito.verify;
 
 import android.annotation.SuppressLint;
+import android.app.admin.ConnectEvent;
+import android.app.admin.DnsEvent;
+import android.app.admin.SecurityLog.SecurityEvent;
 import android.content.Context;
 import android.os.Looper;
 import android.os.PermissionEnforcer;
@@ -40,7 +43,6 @@
 import android.security.forensic.ForensicEvent;
 import android.security.forensic.IForensicServiceCommandCallback;
 import android.security.forensic.IForensicServiceStateCallback;
-import android.util.ArrayMap;
 
 import androidx.test.core.app.ApplicationProvider;
 
@@ -53,7 +55,6 @@
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 
 public class ForensicServiceTest {
     private static final int STATE_UNKNOWN = IForensicServiceStateCallback.State.UNKNOWN;
@@ -300,23 +301,19 @@
         ServiceThread mockThread = spy(ServiceThread.class);
         mDataAggregator.setHandler(mLooperOfDataAggregator, mockThread);
 
-        String eventOneType = "event_one_type";
-        String eventOneMapKey = "event_one_map_key";
-        String eventOneMapVal = "event_one_map_val";
-        Map<String, String> eventOneMap = new ArrayMap<String, String>();
-        eventOneMap.put(eventOneMapKey, eventOneMapVal);
-        ForensicEvent eventOne = new ForensicEvent(eventOneType, eventOneMap);
+        SecurityEvent securityEvent = new SecurityEvent(0, new byte[0]);
+        ForensicEvent eventOne = new ForensicEvent(securityEvent);
 
-        String eventTwoType = "event_two_type";
-        String eventTwoMapKey = "event_two_map_key";
-        String eventTwoMapVal = "event_two_map_val";
-        Map<String, String> eventTwoMap = new ArrayMap<String, String>();
-        eventTwoMap.put(eventTwoMapKey, eventTwoMapVal);
-        ForensicEvent eventTwo = new ForensicEvent(eventTwoType, eventTwoMap);
+        ConnectEvent connectEvent = new ConnectEvent("127.0.0.1", 80, null, 0);
+        ForensicEvent eventTwo = new ForensicEvent(connectEvent);
+
+        DnsEvent dnsEvent = new DnsEvent(null, new String[] {"127.0.0.1"}, 1, null, 0);
+        ForensicEvent eventThree = new ForensicEvent(dnsEvent);
 
         List<ForensicEvent> events = new ArrayList<>();
         events.add(eventOne);
         events.add(eventTwo);
+        events.add(eventThree);
 
         doReturn(true).when(mForensicEventTransportConnection).addData(any());
 
@@ -327,18 +324,16 @@
         ArgumentCaptor<List<ForensicEvent>> captor = ArgumentCaptor.forClass(List.class);
         verify(mForensicEventTransportConnection).addData(captor.capture());
         List<ForensicEvent> receivedEvents = captor.getValue();
-        assertEquals(receivedEvents.size(), 2);
+        assertEquals(receivedEvents.size(), 3);
 
-        assertEquals(receivedEvents.getFirst().getType(), eventOneType);
-        assertEquals(receivedEvents.getFirst().getKeyValuePairs().size(), 1);
-        assertEquals(receivedEvents.getFirst().getKeyValuePairs().get(eventOneMapKey),
-                eventOneMapVal);
+        assertEquals(receivedEvents.get(0).getType(), ForensicEvent.SECURITY_EVENT);
+        assertNotNull(receivedEvents.get(0).getSecurityEvent());
 
-        assertEquals(receivedEvents.getLast().getType(), eventTwoType);
-        assertEquals(receivedEvents.getLast().getKeyValuePairs().size(), 1);
-        assertEquals(receivedEvents.getLast().getKeyValuePairs().get(eventTwoMapKey),
-                eventTwoMapVal);
+        assertEquals(receivedEvents.get(1).getType(), ForensicEvent.NETWORK_EVENT_CONNECT);
+        assertNotNull(receivedEvents.get(1).getConnectEvent());
 
+        assertEquals(receivedEvents.get(2).getType(), ForensicEvent.NETWORK_EVENT_DNS);
+        assertNotNull(receivedEvents.get(2).getDnsEvent());
     }
 
     private class MockInjector implements ForensicService.Injector {