Add a Feature Extractor for the Drakvuf Sandbox

capa/features/address.py

@@ -10,7 +10,8 @@
 
 class Address(abc.ABC):
     @abc.abstractmethod
-    def __eq__(self, other): ...
+    def __eq__(self, other):
+        ...
 
     @abc.abstractmethod
     def __lt__(self, other):

capa/features/common.py

@@ -458,6 +458,7 @@ def evaluate(self, features: "capa.engine.FeatureSet", short_circuit=True):
 FORMAT_SC32 = "sc32"
 FORMAT_SC64 = "sc64"
 FORMAT_CAPE = "cape"
+FORMAT_DRAKVUF = "drakvuf"
 FORMAT_FREEZE = "freeze"
 FORMAT_RESULT = "result"
 STATIC_FORMATS = {
@@ -471,6 +472,7 @@ def evaluate(self, features: "capa.engine.FeatureSet", short_circuit=True):
 }
 DYNAMIC_FORMATS = {
     FORMAT_CAPE,
+    FORMAT_DRAKVUF,
     FORMAT_FREEZE,
     FORMAT_RESULT,
 }

capa/features/extractors/drakvuf/call.py

@@ -0,0 +1,56 @@
+# Copyright (C) 2023 Mandiant, Inc. All Rights Reserved.
+# Licensed 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: [package root]/LICENSE.txt
+# 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.
+
+import logging
+from typing import Tuple, Iterator
+
+from capa.features.insn import API, Number
+from capa.features.common import String, Feature
+from capa.features.address import Address
+from capa.features.extractors.base_extractor import CallHandle, ThreadHandle, ProcessHandle
+from capa.features.extractors.drakvuf.models import Call
+
+logger = logging.getLogger(__name__)
+
+
+def extract_call_features(ph: ProcessHandle, th: ThreadHandle, ch: CallHandle) -> Iterator[Tuple[Feature, Address]]:
+    """
+    this method extracts the given call's features (such as API name and arguments),
+    and returns them as API, Number, and String features.
+
+    args:
+      ph: process handle (for defining the extraction scope)
+      th: thread handle (for defining the extraction scope)
+      ch: call handle (for defining the extraction scope)
+
+    yields:
+      Feature, address; where Feature is either: API, Number, or String.
+    """
+    call: Call = ch.inner
+
+    # list similar to disassembly: arguments right-to-left, call
+    for arg_value in call.arguments.values():
+        if arg_value.isdecimal():
+            yield Number(int(arg_value)), ch.address
+        else:
+            try:
+                yield Number(int(arg_value, 16)), ch.address
+            except:
+                # yield it as a string
+                yield String(arg_value), ch.address
+
+    yield API(call.name), ch.address
+
+
+def extract_features(ph: ProcessHandle, th: ThreadHandle, ch: CallHandle) -> Iterator[Tuple[Feature, Address]]:
+    for handler in CALL_HANDLERS:
+        for feature, addr in handler(ph, th, ch):
+            yield feature, addr
+
+
+CALL_HANDLERS = (extract_call_features,)

capa/features/extractors/drakvuf/extractor.py

@@ -0,0 +1,98 @@
+# Copyright (C) 2023 Mandiant, Inc. All Rights Reserved.
+# Licensed 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: [package root]/LICENSE.txt
+# 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.
+
+import logging
+from typing import Dict, List, Tuple, Union, Iterator
+
+import capa.features.extractors.drakvuf.call
+import capa.features.extractors.drakvuf.file
+import capa.features.extractors.drakvuf.thread
+import capa.features.extractors.drakvuf.global_
+import capa.features.extractors.drakvuf.process
+from capa.features.common import Feature, Characteristic
+from capa.features.address import NO_ADDRESS, Address, ThreadAddress, ProcessAddress, AbsoluteVirtualAddress, _NoAddress
+from capa.features.extractors.base_extractor import (
+    CallHandle,
+    SampleHashes,
+    ThreadHandle,
+    ProcessHandle,
+    DynamicFeatureExtractor,
+)
+from capa.features.extractors.drakvuf.models import Call, DrakvufReport
+from capa.features.extractors.drakvuf.helpers import sort_calls
+
+logger = logging.getLogger(__name__)
+
+
+class DrakvufExtractor(DynamicFeatureExtractor):
+    def __init__(self, report: DrakvufReport):
+        super().__init__(
+            # DRAKVUF currently does not yield hash information about the sample in its output
+            hashes=SampleHashes(md5="", sha1="", sha256="")
+        )
+
+        self.report: DrakvufReport = report
+
+        # sort the api calls to prevent going through the entire list each time
+        self.sorted_calls: Dict[ProcessAddress, Dict[ThreadAddress, Call]] = sort_calls(report)
+
+        # pre-compute these because we'll yield them at *every* scope.
+        self.global_features = list(capa.features.extractors.drakvuf.global_.extract_features(self.report))
+
+    def get_base_address(self) -> Union[AbsoluteVirtualAddress, _NoAddress, None]:
+        # DRAKVUF currently does not yield information about the PE's address
+        return _NoAddress
+
+    def extract_global_features(self) -> Iterator[Tuple[Feature, Address]]:
+        yield from self.global_features
+
+    def extract_file_features(self) -> Iterator[Tuple[Feature, Address]]:
+        yield from capa.features.extractors.drakvuf.file.extract_features(self.report)
+
+    def get_processes(self) -> Iterator[ProcessHandle]:
+        yield from capa.features.extractors.drakvuf.file.get_processes(self.sorted_calls)
+
+    def extract_process_features(self, ph: ProcessHandle) -> Iterator[Tuple[Feature, Address]]:
+        yield from capa.features.extractors.drakvuf.process.extract_features(ph)
+
+    def get_process_name(self, ph: ProcessHandle) -> str:
+        return ph.inner["process_name"]
+
+    def get_threads(self, ph: ProcessHandle) -> Iterator[ThreadHandle]:
+        yield from capa.features.extractors.drakvuf.process.get_threads(self.sorted_calls, ph)
+
+    def extract_thread_features(self, ph: ProcessHandle, th: ThreadHandle) -> Iterator[Tuple[Feature, Address]]:
+        if False:
+            # force this routine to be a generator,
+            # but we don't actually have any elements to generate.
+            yield Characteristic("never"), NO_ADDRESS
+        return
+
+    def get_calls(self, ph: ProcessHandle, th: ThreadHandle) -> Iterator[CallHandle]:
+        yield from capa.features.extractors.drakvuf.thread.get_calls(self.sorted_calls, ph, th)
+
+    def get_call_name(self, ph: ProcessHandle, th: ThreadHandle, ch: CallHandle) -> str:
+        call: Call = ch.inner
+        call_name = "{}({}){}".format(
+            call.name,
+            ", ".join(f"{arg_name}={arg_value}" for arg_name, arg_value in call.arguments.items()),
+            f" -> {call.return_value}"
+            if hasattr(call, "return_value")
+            else "",  # SysCalls don't have a return value, while WinApi calls do
+        )
+        return call_name
+
+    def extract_call_features(
+        self, ph: ProcessHandle, th: ThreadHandle, ch: CallHandle
+    ) -> Iterator[Tuple[Feature, Address]]:
+        yield from capa.features.extractors.drakvuf.call.extract_features(ph, th, ch)
+
+    @classmethod
+    def from_report(cls, report: List[Dict]) -> "DrakvufExtractor":
+        dr = DrakvufReport.from_raw_report(report)
+        return DrakvufExtractor(report=dr)

capa/features/extractors/drakvuf/file.py

@@ -0,0 +1,62 @@
+# Copyright (C) 2023 Mandiant, Inc. All Rights Reserved.
+# Licensed 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: [package root]/LICENSE.txt
+# 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.
+
+import logging
+from typing import Dict, Tuple, Iterator
+
+from capa.features.file import Import
+from capa.features.common import String, Feature
+from capa.features.address import NO_ADDRESS, Address, ThreadAddress, ProcessAddress, AbsoluteVirtualAddress
+from capa.features.extractors.helpers import generate_symbols
+from capa.features.extractors.base_extractor import ProcessHandle
+from capa.features.extractors.drakvuf.models import Call, DrakvufReport
+
+logger = logging.getLogger(__name__)
+
+
+def get_processes(calls: Dict[ProcessAddress, Dict[ThreadAddress, Call]]) -> Iterator[ProcessHandle]:
+    """
+    get all the created processes for a sample
+    """
+    for proc_addr, calls_per_thread in calls.items():
+        sample_call: Call = next(iter(calls_per_thread.values()))[0]  # get process name
+        yield ProcessHandle(proc_addr, inner={"process_name": sample_call.process_name})
+
+
+def extract_import_names(report: DrakvufReport) -> Iterator[Tuple[Feature, Address]]:
+    """
+    extract imported function names
+    """
+    if report.loaded_dlls is None:
+        return
+    dlls = report.loaded_dlls
+
+    for dll in dlls:
+        dll_base_name = dll.name.split("\\")[-1]
+        for function_name, function_address in dll.imports.items():
+            for name in generate_symbols(dll_base_name, function_name, include_dll=True):
+                yield Import(name), AbsoluteVirtualAddress(function_address)
+
+
+def extract_file_strings(report: DrakvufReport) -> Iterator[Tuple[Feature, Address]]:
+    if report.discovered_dlls is None:
+        return
+    for dll in report.discovered_dlls:
+        yield String(dll.name), NO_ADDRESS
+
+
+def extract_features(report: DrakvufReport) -> Iterator[Tuple[Feature, Address]]:
+    for handler in FILE_HANDLERS:
+        for feature, addr in handler(report):
+            yield feature, addr
+
+
+FILE_HANDLERS = (
+    extract_import_names,
+    extract_file_strings,
+)

capa/features/extractors/drakvuf/global_.py

@@ -0,0 +1,38 @@
+# Copyright (C) 2023 Mandiant, Inc. All Rights Reserved.
+# Licensed 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: [package root]/LICENSE.txt
+# 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.
+
+import logging
+from typing import Tuple, Iterator
+
+from capa.features.common import OS, FORMAT_PE, OS_WINDOWS, Format, Feature
+from capa.features.address import NO_ADDRESS, Address
+from capa.features.extractors.drakvuf.models import DrakvufReport
+
+logger = logging.getLogger(__name__)
+
+
+def extract_format(report: DrakvufReport) -> Iterator[Tuple[Feature, Address]]:
+    # drakvuf sandbox currently supports only windows as the guest: https://drakvuf-sandbox.readthedocs.io/en/latest/usage/getting_started.html
+    yield Format(FORMAT_PE), NO_ADDRESS
+
+
+def extract_os(report: DrakvufReport) -> Iterator[Tuple[Feature, Address]]:
+    # drakvuf sandbox currently supports only windows as the guest: https://drakvuf-sandbox.readthedocs.io/en/latest/usage/getting_started.html
+    yield OS(OS_WINDOWS), NO_ADDRESS
+
+
+def extract_features(report: DrakvufReport) -> Iterator[Tuple[Feature, Address]]:
+    for global_handler in GLOBAL_HANDLER:
+        for feature, addr in global_handler(report):
+            yield feature, addr
+
+
+GLOBAL_HANDLER = (
+    extract_format,
+    extract_os,
+)

capa/features/extractors/drakvuf/helpers.py

@@ -0,0 +1,34 @@
+# Copyright (C) 2023 Mandiant, Inc. All Rights Reserved.
+# Licensed 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: [package root]/LICENSE.txt
+# 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.
+
+from typing import Dict
+
+from capa.features.address import ThreadAddress, ProcessAddress
+from capa.features.extractors.drakvuf.models import Call, DrakvufReport
+
+
+def sort_calls(report: DrakvufReport) -> Dict[ProcessAddress, Dict[ThreadAddress, Call]]:
+    result = {}
+    for call in (*report.syscalls, *report.apicalls):
+        if call.pid < 1:
+            # ignore pid's with zero
+            continue
+        proc_addr = ProcessAddress(pid=call.pid, ppid=call.ppid)
+        thread_addr = ThreadAddress(process=proc_addr, tid=call.tid)
+        if proc_addr not in result:
+            result[proc_addr] = {}
+        if thread_addr not in result[proc_addr]:
+            result[proc_addr][thread_addr] = []
+
+        result[proc_addr][thread_addr].append(call)
+
+    for proc, threads in result.items():
+        for thread in threads:
+            result[proc][thread].sort(key=lambda call: call.timestamp)
+
+    return result