dotnet: emit calls to/from MethodDef methods (#1236)

* dotnet: emit calls to/from MethodDef methods

* dotnet: update function.py copyright header

CHANGELOG.md

@@ -8,6 +8,7 @@
 - dotnet: emit API features for objects created via the newobj instruction #1186 @mike-hunhoff
 - dotnet: emit API features for generic methods #1231 @mike-hunhoff
 - Python 3.11 support #1192 @williballenthin
+- dotnet: emit calls to/from MethodDef methods #1236 @mike-hunhoff
 
 ### Breaking Changes
 

capa/features/extractors/dnfile/extractor.py

@@ -8,13 +8,15 @@
 
 from __future__ import annotations
 
-from typing import List, Tuple, Iterator
+from typing import Dict, List, Tuple, Iterator, Optional
 
 import dnfile
+from dncil.cil.opcode import OpCodes
 
 import capa.features.extractors
 import capa.features.extractors.dnfile.file
 import capa.features.extractors.dnfile.insn
+import capa.features.extractors.dnfile.function
 from capa.features.common import Feature
 from capa.features.address import NO_ADDRESS, Address, DNTokenAddress, DNTokenOffsetAddress
 from capa.features.extractors.base_extractor import BBHandle, InsnHandle, FunctionHandle, FeatureExtractor
@@ -41,12 +43,36 @@ class DnfileFeatureExtractor(FeatureExtractor):
         yield from capa.features.extractors.dnfile.file.extract_features(self.pe)
 
     def get_functions(self) -> Iterator[FunctionHandle]:
-        for token, f in get_dotnet_managed_method_bodies(self.pe):
-            yield FunctionHandle(address=DNTokenAddress(token), inner=f, ctx={"pe": self.pe})
+        # create a method lookup table
+        methods: Dict[Address, FunctionHandle] = {}
+        for (token, method) in get_dotnet_managed_method_bodies(self.pe):
+            fh: FunctionHandle = FunctionHandle(
+                address=DNTokenAddress(token), inner=method, ctx={"pe": self.pe, "calls_from": set(), "calls_to": set()}
+            )
 
-    def extract_function_features(self, f):
-        # TODO
-        yield from []
+            # method tokens should be unique
+            assert fh.address not in methods.keys()
+            methods[fh.address] = fh
+
+        # calculate unique calls to/from each method
+        for fh in methods.values():
+            for insn in fh.inner.instructions:
+                if insn.opcode not in (OpCodes.Call, OpCodes.Callvirt, OpCodes.Jmp, OpCodes.Calli, OpCodes.Newobj):
+                    continue
+
+                # record call to destination method; note: we only consider MethodDef methods for destinations
+                dest: Optional[FunctionHandle] = methods.get(DNTokenAddress(insn.operand.value), None)
+                if dest is not None:
+                    dest.ctx["calls_to"].add(fh.address)
+
+                # record call from source method; note: we record all unique calls from a MethodDef method, not just
+                # those calls to other MethodDef methods e.g. calls to imported MemberRef methods
+                fh.ctx["calls_from"].add(DNTokenAddress(insn.operand.value))
+
+        yield from methods.values()
+
+    def extract_function_features(self, fh) -> Iterator[Tuple[Feature, Address]]:
+        yield from capa.features.extractors.dnfile.function.extract_features(fh)
 
     def get_basic_blocks(self, f) -> Iterator[BBHandle]:
         # each dotnet method is considered 1 basic block

capa/features/extractors/dnfile/function.py

@@ -0,0 +1,50 @@
+# Copyright (C) 2020 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 __future__ import annotations
+
+import logging
+from typing import Tuple, Iterator
+
+from capa.features.common import Feature, Characteristic
+from capa.features.address import Address
+from capa.features.extractors.base_extractor import FunctionHandle
+
+logger = logging.getLogger(__name__)
+
+
+def extract_function_calls_to(fh: FunctionHandle) -> Iterator[Tuple[Characteristic, Address]]:
+    """extract callers to a function"""
+    for dest in fh.ctx["calls_to"]:
+        yield Characteristic("calls to"), dest
+
+
+def extract_function_calls_from(fh: FunctionHandle) -> Iterator[Tuple[Characteristic, Address]]:
+    """extract callers from a function"""
+    for src in fh.ctx["calls_from"]:
+        yield Characteristic("calls from"), src
+
+
+def extract_recursive_call(fh: FunctionHandle) -> Iterator[Tuple[Characteristic, Address]]:
+    """extract recursive function call"""
+    if fh.address in fh.ctx["calls_to"]:
+        yield Characteristic("recursive call"), fh.address
+
+
+def extract_function_loop(fh: FunctionHandle) -> Iterator[Tuple[Characteristic, Address]]:
+    """extract loop indicators from a function"""
+    raise NotImplementedError()
+
+
+def extract_features(fh: FunctionHandle) -> Iterator[Tuple[Feature, Address]]:
+    for func_handler in FUNCTION_HANDLERS:
+        for (feature, addr) in func_handler(fh):
+            yield feature, addr
+
+
+FUNCTION_HANDLERS = (extract_function_calls_to, extract_function_calls_from, extract_recursive_call)

tests/fixtures.py

@@ -753,6 +753,10 @@ FEATURE_PRESENCE_TESTS_DOTNET = sorted(
         ("_1c444", "function=0x1F68", capa.features.insn.API("GetWindowDC"), True),
         ("_1c444", "function=0x1F68", capa.features.insn.API("user32.GetWindowDC"), True),
         ("_1c444", "function=0x1F68", capa.features.insn.Number(0xCC0020), True),
+        ("_1c444", "token=0x600001D", capa.features.common.Characteristic("calls to"), True),
+        ("_1c444", "token=0x6000018", capa.features.common.Characteristic("calls to"), False),
+        ("_1c444", "token=0x600001D", capa.features.common.Characteristic("calls from"), True),
+        ("_1c444", "token=0x600000F", capa.features.common.Characteristic("calls from"), False),
         ("_1c444", "function=0x1F68", capa.features.insn.Number(0x0), True),
         ("_1c444", "function=0x1F68", capa.features.insn.Number(0x1), False),
         ("_692f", "token=0x6000004", capa.features.insn.API("System.Linq.Enumerable::First"), True),  # generic method
@@ -950,7 +954,10 @@ FEATURE_COUNT_TESTS = [
 ]
 
 
-FEATURE_COUNT_TESTS_DOTNET = []  # type: ignore
+FEATURE_COUNT_TESTS_DOTNET = [
+    ("_1c444", "token=0x600001D", capa.features.common.Characteristic("calls to"), 1),
+    ("_1c444", "token=0x600001D", capa.features.common.Characteristic("calls from"), 9),
+]
 
 
 def do_test_feature_presence(get_extractor, sample, scope, feature, expected):