pyinstaller: fix viv pe parser

release v1.3.0

.github/ISSUE_TEMPLATE/bug_report.md

@@ -3,26 +3,29 @@ name: Bug report
 about: Create a report to help us improve
 
 ---
-
 <!--
+# Is your bug report related to capa rules (for example a false positive)?
+We use sybmodules to separate code, rules and test data. If your issue is related to capa rules, please report it at https://github.com/fireeye/capa-rules/issues.
 
-Have you read capa's Code of Conduct? By filing an Issue, you are expected to comply with it, including treating everyone with respect: https://github.com/fireeye/capa/blob/master/.github/CODE_OF_CONDUCT.md
+# Have you checked that your issue isn't already filed?
+Please search if there is a similar issue at https://github.com/fireeye/capa/issues. If there is already a similar issue, please add more details there instead of opening a new one.
+
+# Have you read capa's Code of Conduct?
+By filing an Issue, you are expected to comply with it, including treating everyone with respect: https://github.com/fireeye/capa/blob/master/.github/CODE_OF_CONDUCT.md
+
+# Have you read capa's CONTRIBUTING guide?
+It contains helpful information about how to contribute to capa. Check https://github.com/fireeye/capa/blob/master/.github/CONTRIBUTING.md#reporting-bugs
 -->
 
-### Prerequisites
-
-* [ ] Put an X between the brackets on this line if you have done all of the following:
-    * Checked that your issue isn't already filed: [search](https://github.com/fireeye/capa/issues?q=is%3Aissue+is%3Aopen+)
-
 ### Description
 
 <!-- Description of the issue -->
 
 ### Steps to Reproduce
 
-1. <!-- First Step -->
-2. <!-- Second Step -->
-3. <!-- and so on… -->
+<!-- 1. First Step -->
+<!-- 2. Second Step -->
+<!-- 3. and so on… -->
 
 **Expected behavior:**
 
@@ -32,10 +35,6 @@ Have you read capa's Code of Conduct? By filing an Issue, you are expected to co
 
 <!-- What actually happens -->
 
-**Reproduces how often:**
-
-<!-- What percentage of the time does it reproduce? -->
-
 ### Versions
 
 <!-- You can get this information from copy and pasting the output of `capa --version` from the command line.
@@ -45,3 +44,4 @@ Have you read capa's Code of Conduct? By filing an Issue, you are expected to co
 ### Additional Information
 
 <!-- Any additional information, configuration or data that might be necessary to reproduce the issue. -->
+

.github/ISSUE_TEMPLATE/feature_request.md

@@ -3,24 +3,33 @@ name: Feature request
 about: Suggest an idea for capa
 
 ---
-
 <!--
+# Is your issue related to capa rules (for example an idea for a new rule)?
+We use sybmodules to separate code, rules and test data. If your issue is related to capa rules, please report it at https://github.com/fireeye/capa-rules/issues.
 
-Have you read capa's Code of Conduct? By filing an Issue, you are expected to comply with it, including treating everyone with respect: https://github.com/fireeye/capa/blob/master/CODE_OF_CONDUCT.md
+# Have you checked that your issue isn't already filed?
+Please search if there is a similar issue at https://github.com/fireeye/capa/issues. If there is already a similar issue, please add more details there instead of opening a new one.
+
+# Have you read capa's Code of Conduct?
+By filing an Issue, you are expected to comply with it, including treating everyone with respect: https://github.com/fireeye/capa/blob/master/.github/CODE_OF_CONDUCT.md
+
+# Have you read capa's CONTRIBUTING guide?
+It contains helpful information about how to contribute to capa. Check https://github.com/fireeye/capa/blob/master/.github/CONTRIBUTING.md#suggesting-enhancements
 -->
 
-## Summary
+### Summary
 
 <!-- One paragraph explanation of the feature. -->
 
-## Motivation
+### Motivation
 
 <!-- Why are we doing this? What use cases does it support? What is the expected outcome? -->
 
-## Describe alternatives you've considered
+### Describe alternatives you've considered
 
 <!-- A clear and concise description of the alternative solutions you've considered. -->
 
 ## Additional context
 
 <!-- Add any other context or screenshots about the feature request here. -->
+

.github/pyinstaller/pyinstaller.spec

@@ -9,8 +9,16 @@ import wcwidth
 # when invoking pyinstaller from the project root,
 # this gets run from the project root.
 with open('./capa/version.py', 'wb') as f:
-    f.write("__version__ = '%s'"
-            % subprocess.check_output(["git", "describe", "--always"]).strip())
+    # git output will look like:
+    #
+    #     tags/v1.0.0-0-g3af38dc
+    #         ------- tag
+    #                 - commits since
+    #                   g------- git hash fragment
+    version = (subprocess.check_output(["git", "describe", "--always", "--tags", "--long"])
+               .strip()
+               .replace("tags/", ""))
+    f.write("__version__ = '%s'" % version)
 
 a = Analysis(
     # when invoking pyinstaller from the project root,
@@ -36,7 +44,6 @@ a = Analysis(
     hiddenimports=[
         # vivisect does manual/runtime importing of its modules,
         # so declare the things that could be imported here.
-        "pycparser",
         "vivisect",
         "vivisect.analysis",
         "vivisect.analysis.amd64",
@@ -84,11 +91,13 @@ a = Analysis(
         "vivisect.impapi.windows",
         "vivisect.impapi.windows.amd64",
         "vivisect.impapi.windows.i386",
+        "vivisect.impapi.winkern.i386",
+        "vivisect.impapi.winkern.amd64",
         "vivisect.parsers.blob",
         "vivisect.parsers.elf",
         "vivisect.parsers.ihex",
         "vivisect.parsers.macho",
-        "vivisect.parsers.parse_pe",
+        "vivisect.parsers.pe",
         "vivisect.parsers.utils",
         "vivisect.storage",
         "vivisect.storage.basicfile",

.github/workflows/build.yml

@@ -0,0 +1,77 @@
+name: build
+
+on:
+  release:
+    types: [created, edited]
+
+jobs:
+  build:
+    name: PyInstaller for ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        include:
+          - os: ubuntu-16.04
+            # use old linux so that the shared library versioning is more portable
+            artifact_name: capa
+            asset_name: linux
+          - os: windows-latest
+            artifact_name: capa.exe
+            asset_name: windows
+          - os: macos-latest
+            artifact_name: capa
+            asset_name: macos
+    steps:
+      - name: Checkout capa
+        uses: actions/checkout@v2
+        with:
+          submodules: true
+      - name: Set up Python 2.7
+        uses: actions/setup-python@v2
+        with:
+          python-version: 2.7
+      - name: Install PyInstaller
+        # pyinstaller 4 doesn't support Python 2.7
+        run: pip install 'pyinstaller==3.*'
+      - name: Install capa
+        run: pip install -e .
+      - name: Build standalone executable
+        run: pyinstaller .github/pyinstaller/pyinstaller.spec
+      - name: Does it run?
+        run: dist/capa "tests/data/Practical Malware Analysis Lab 01-01.dll_"
+      - uses: actions/upload-artifact@v2
+        with:
+          name: ${{ matrix.asset_name }}
+          path: dist/${{ matrix.artifact_name }}
+
+  zip:
+    name: zip ${{ matrix.asset_name }}
+    runs-on: ubuntu-latest
+    needs: build
+    strategy:
+      matrix:
+        include:
+          - asset_name: linux
+            artifact_name: capa
+          - asset_name: windows
+            artifact_name: capa.exe
+          - asset_name: macos
+            artifact_name: capa
+    steps:
+      - name: Download ${{ matrix.asset_name }}
+        uses: actions/download-artifact@v2
+        with:
+          name: ${{ matrix.asset_name }}
+      - name: Set executable flag
+        run: chmod +x ${{ matrix.artifact_name }}
+      - name: Set zip name
+        run: echo ::set-env name=zip_name::capa-${GITHUB_REF#refs/tags/}-${{ matrix.asset_name }}.zip
+      - name: Zip ${{ matrix.artifact_name }} into ${{ env.zip_name }}
+        run: zip ${{ env.zip_name }} ${{ matrix.artifact_name }}
+      - name: Upload ${{ env.zip_name }} to GH Release
+        uses: svenstaro/upload-release-action@v2
+        with:
+          repo_token: ${{ secrets.GITHUB_TOKEN}}
+          file: ${{ env.zip_name }}
+          tag: ${{ github.ref }}
+

.github/workflows/publish.yml

@@ -0,0 +1,29 @@
+# This workflows will upload a Python Package using Twine when a release is created
+# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries
+
+name: publish to pypi
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  deploy:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: '2.7'
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install setuptools wheel twine
+      - name: Build and publish
+        env:
+          TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
+          TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
+        run: |
+          python setup.py sdist bdist_wheel
+          twine upload --skip-existing dist/*

.github/workflows/tests.yml

@@ -0,0 +1,68 @@
+name: CI
+
+on:
+  push:
+    branches: [ master ]
+  pull_request:
+    branches: [ master ]
+
+jobs:
+  code_style:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout capa
+      uses: actions/checkout@v2
+    - name: Set up Python 3.8
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.8
+    - name: Install dependencies
+      run: pip install 'isort==5.*' black
+    - name: Lint with isort
+      run: isort --profile black --length-sort --line-width 120 -c .
+    - name: Lint with black
+      run: black -l 120 --check .
+
+  rule_linter:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Checkout capa with rules submodule
+      uses: actions/checkout@v2
+      with:
+        submodules: true
+    - name: Set up Python 3.8
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.8
+    # We don't need vivisect, so we can install capa using Python3
+    - name: Install capa
+      run: pip install -e .
+    - name: Run rule linter
+      run: python scripts/lint.py rules/
+
+  tests:
+    name: Tests in ${{ matrix.python }}
+    runs-on: ubuntu-latest
+    needs: [code_style, rule_linter]
+    strategy:
+      matrix:
+        include:
+          - python: 2.7
+          - python: 3.6
+          - python: 3.7
+          - python: 3.8
+          - python: '3.9.0-rc.1' # Python latest
+    steps:
+    - name: Checkout capa with submodules
+      uses: actions/checkout@v2
+      with:
+        submodules: true
+    - name: Set up Python ${{ matrix.python }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python }}
+    - name: Install capa
+      run: pip install -e .[dev]
+    - name: Run tests
+      run: pytest tests/
+

.gitignore

@@ -110,6 +110,7 @@ venv.bak/
 *.i64
 !rules/lib
 
-# hooks output
-style-checker-output.log
+# hooks/ci.sh output
+isort-output.log
+black-output.log
 rule-linter-output.log

.gitmodules

@@ -1,6 +1,6 @@
 [submodule "rules"]
 	path = rules
-	url = git@github.com:fireeye/capa-rules.git
+	url = ../capa-rules.git
 [submodule "tests/data"]
 	path = tests/data
-	url = git@github.com:fireeye/capa-testfiles.git
+	url = ../capa-testfiles.git

CHANGELOG.md

@@ -0,0 +1,258 @@
+# Change Log
+
+## v1.3.0 (2020-09-14)
+
+This release brings newly updated mappings to the [Malware Behavior Catalog version 2.0](https://github.com/MBCProject/mbc-markdown), many enhancements to the IDA Pro plugin, [flare-capa on PyPI](https://pypi.org/project/flare-capa/), a bunch of bug fixes to improve feature extraction, and four new rules. We received contributions from ten reverse engineers, including seven new ones:
+
+  - @dzbeck
+  - @recvfrom
+  - @toomanybananas
+  - @cclauss 
+  - @adamprescott91 
+  - @weslambert
+  - @stevemk14ebr 
+  
+Download a standalone binary below and checkout the readme [here on GitHub](https://github.com/fireeye/capa/). Report issues on our [issue tracker](https://github.com/fireeye/capa/issues) and contribute new rules at [capa-rules](https://github.com/fireeye/capa-rules/).
+
+### Key changes to IDA Plugin
+
+The IDA Pro integration is now distributed as a real plugin, instead of a script. This enables a few things:
+
+  - keyboard shortcuts and file menu integration
+  - updates distributed PyPI/`pip install --upgrade` without touching your `%IDADIR%`
+  - generally doing thing the "right way"
+
+How to get this new version? Its easy: download [capa_explorer.py](https://raw.githubusercontent.com/fireeye/capa/master/capa/ida/plugin/capa_explorer.py) to your IDA plugins directory and update your capa installation (incidentally, this is a good opportunity to migrate to `pip install flare-capa` instead of git checkouts). Now you should see the plugin listed in the `Edit > Plugins > FLARE capa explorer` menu in IDA. 
+
+Please refer to the plugin [readme](https://github.com/fireeye/capa/blob/master/capa/ida/plugin/README.md) for additional information on installing and using the IDA Pro plugin.
+
+Please open an issue in this repository if you notice anything weird.
+ 
+### New features
+
+  - ida plugin: now a real plugin, not a script @mike-hunhoff 
+  - core: distributed via PyPI as [flare-capa](https://pypi.org/project/flare-capa/) @williballenthin 
+  - features: enable automatic A/W handling for imports @williballenthin @Ana06 #246 
+  - ida plugin: persist rules directory setting via [ida-settings](https://github.com/williballenthin/ida-settings) @williballenthin #268
+  - ida plugin: add search bar to results view @williballenthin #285
+  - ida plugin: add `Analyze` and `Reset` buttons to tree view @mike-hunhoff #304
+  - ida plugin: add status label to tree view @mike-hunhoff
+  - ida plugin: add progress indicator @mike-hunhoff, @mr-tz
+
+### New rules
+
+  - compiled with py2exe @re-fox
+  - resolve path using msvcrt @re-fox 
+  - decompress data using QuickLZ @edeca
+  - encrypt data using sosemanuk @recvfrom 
+
+### Bug fixes
+
+  - rule: reduce FP in DNS resolution @toomanybananas
+  - engine: report correct strings matched via regex @williballenthin #262 
+  - formatter: correctly format descriptions in two-line syntax @williballenthin @recvfrom #263 
+  - viv: better extract offsets from SibOper operands @williballenthin @edeca #276 
+  - import-to-ida: fix import error @cclauss 
+  - viv: don't write settings to ~/.viv/viv.json @williballenthin @rakuy0 @weslambert #244
+  - ida plugin: remove dependency loop that resulted in unnecessary overhead @mike-hunhoff #303
+  - ida plugin: correctly highlight regex matches in IDA Disassembly view @mike-hunhoff #305
+  - ida plugin: better handle rule directory prompt and failure case @stevemk14ebr @mike-hunhoff #309
+
+### Changes
+
+  - rules: update meta mapping to MBC 2.0! @dzbeck
+  - render: don't display rules that are also matched by other rules @williballenthin @Ana06 #224
+  - ida plugin: simplify tabs, removing summary and adding detail to results view @williballenthin #286
+  - ida plugin: analysis is no longer automatically started when plugin is first opened @mike-hunhoff #304
+  - ida plugin: user must manually select a capa rules directory before analysis can be performed @mike-hunhoff
+  - ida plugin: user interface controls are disabled until analysis is performed @mike-hunhoff #304
+
+### Raw diffs
+
+  - [capa v1.2.0...v1.3.0](https://github.com/fireeye/capa/compare/v1.2.0...v1.3.0)
+  - [capa-rules v1.2.0...v1.3.0](https://github.com/fireeye/capa-rules/compare/v1.2.0...v1.3.0)
+
+## v1.2.0 (2020-08-31)
+
+This release brings UI enhancements, especially for the IDA Pro plugin, 
+investment towards py3 support,
+fixes some bugs identified by the community, 
+and 46 (!) new rules.
+We received contributions from ten reverse engineers, including five new ones:
+
+  - @agithubuserlol
+  - @recvfrom
+  - @D4nch3n
+  - @edeca
+  - @winniepe 
+  
+Download a standalone binary below and checkout the readme [here on GitHub](https://github.com/fireeye/capa/).
+Report issues on our [issue tracker](https://github.com/fireeye/capa/issues)
+and contribute new rules at [capa-rules](https://github.com/fireeye/capa-rules/).
+ 
+### New features
+
+  - ida plugin: display arch flavors @mike-hunhoff
+  - ida plugin: display block descriptions @mike-hunhoff
+  - ida backend: extract features from nested pointers @mike-hunhoff
+  - main: show more progress output @williballenthin
+  - core: pin dependency versions #258 @recvfrom
+
+### New rules
+  - bypass UAC via AppInfo ALPC @agithubuserlol
+  - bypass UAC via token manipulation @agithubuserlol
+  - check for sandbox and av modules @re-fox
+  - check for sandbox username @re-fox
+  - check if process is running under wine @re-fox
+  - validate credit card number using luhn algorithm @re-fox
+  - validate credit card number using luhn algorithm with no lookup table @re-fox
+  - hash data using FNV @edeca @mr-tz
+  - link many functions at runtime @mr-tz
+  - reference public RSA key @mr-tz
+  - packed with ASPack @williballenthin
+  - delete internet cache @mike-hunhoff
+  - enumerate internet cache @mike-hunhoff
+  - send ICMP echo request @mike-hunhoff
+  - check for debugger via API @mike-hunhoff
+  - check for hardware breakpoints @mike-hunhoff
+  - check for kernel debugger via shared user data structure @mike-hunhoff
+  - check for protected handle exception @mike-hunhoff
+  - check for software breakpoints @mike-hunhoff
+  - check for trap flag exception @mike-hunhoff
+  - check for unexpected memory writes @mike-hunhoff
+  - check process job object @mike-hunhoff
+  - reference anti-VM strings targeting Parallels @mike-hunhoff
+  - reference anti-VM strings targeting Qemu @mike-hunhoff
+  - reference anti-VM strings targeting VirtualBox @mike-hunhoff
+  - reference anti-VM strings targeting VirtualPC @mike-hunhoff
+  - reference anti-VM strings targeting VMWare @mike-hunhoff
+  - reference anti-VM strings targeting Xen @mike-hunhoff
+  - reference analysis tools strings @mike-hunhoff
+  - reference WMI statements @mike-hunhoff
+  - get number of processor cores @mike-hunhoff
+  - get number of processors @mike-hunhoff
+  - enumerate disk properties @mike-hunhoff
+  - get disk size @mike-hunhoff
+  - get process heap flags @mike-hunhoff
+  - get process heap force flags @mike-hunhoff
+  - get Explorer PID @mike-hunhoff
+  - delay execution @mike-hunhoff
+  - check for process debug object @mike-hunhoff
+  - check license value @mike-hunhoff
+  - check ProcessDebugFlags @mike-hunhoff
+  - check ProcessDebugPort @mike-hunhoff
+  - check SystemKernelDebuggerInformation @mike-hunhoff
+  - check thread yield allowed @mike-hunhoff
+  - enumerate system firmware tables @mike-hunhoff
+  - get system firmware table @mike-hunhoff
+  - hide thread from debugger @mike-hunhoff
+
+### Bug fixes
+
+  - ida backend: extract unmapped immediate number features @mike-hunhoff
+  - ida backend: fix stack cookie check #257 @mike-hunhoff
+  - viv backend: better extract gs segment access @williballenthin
+  - core: enable counting of string features #241 @D4nch3n @williballenthin
+  - core: enable descriptions on feature with arch flavors @mike-hunhoff
+  - core: update git links for non-SSH access #259 @recvfrom
+
+### Changes
+
+  - ida plugin: better default display showing first level nesting @winniepe
+  - remove unused `characteristic(switch)` feature @ana06
+  - prepare testing infrastructure for multiple backends/py3 @williballenthin
+  - ci: zip build artifacts @ana06
+  - ci: build all supported python versions @ana06
+  - code style and formatting @mr-tz
+
+### Raw diffs
+
+  - [capa v1.1.0...v1.2.0](https://github.com/fireeye/capa/compare/v1.1.0...v1.2.0)
+  - [capa-rules v1.1.0...v1.2.0](https://github.com/fireeye/capa-rules/compare/v1.1.0...v1.2.0)
+
+## v1.1.0 (2020-08-05)
+
+This release brings new rule format updates, such as adding `offset/x32` and negative offsets,
+fixes some bugs identified by the community, and 28 (!) new rules.
+We received contributions from eight reverse engineers, including four new ones:
+
+  - @re-fox
+  - @psifertex
+  - @bitsofbinary
+  - @threathive
+  
+Download a standalone binary below and checkout the readme [here on GitHub](https://github.com/fireeye/capa/). Report issues on our [issue tracker](https://github.com/fireeye/capa/issues) and contribute new rules at [capa-rules](https://github.com/fireeye/capa-rules/).
+  
+### New features
+
+  - import: add Binary Ninja import script #205 #207 @psifertex
+  - rules: offsets can be negative #197 #208 @williballenthin
+  - rules: enable descriptions for statement nodes #194 #209 @Ana06
+  - rules: add arch flavors to number and offset features #210 #216 @williballenthin
+  - render: show SHA1/SHA256 in default report #164 @threathive
+  - tests: add tests for IDA Pro backend #202 @williballenthin
+  
+### New rules
+
+  - check for unmoving mouse cursor @BitsOfBinary
+  - check mutex and exit @re-fox
+  - parse credit card information @re-fox
+  - read ini file @re-fox
+  - validate credit card number with luhn algorithm @re-fox
+  - change the wallpaper @re-fox
+  - acquire debug privileges @williballenthin
+  - import public key @williballenthin
+  - terminate process by name @williballenthin
+  - encrypt data using DES @re-fox
+  - encrypt data using DES via WinAPI @re-fox
+  - hash data using sha1 via x86 extensions @re-fox
+  - hash data using sha256 via x86 extensions @re-fox
+  - capture network configuration via ipconfig @re-fox
+  - hash data via WinCrypt @mike-hunhoff
+  - get file attributes @mike-hunhoff
+  - allocate thread local storage @mike-hunhoff
+  - get thread local storage value @mike-hunhoff
+  - set thread local storage @mike-hunhoff
+  - get session integrity level @mike-hunhoff
+  - add file to cabinet file @mike-hunhoff
+  - flush cabinet file @mike-hunhoff
+  - open cabinet file @mike-hunhoff
+  - gather firefox profile information @re-fox
+  - encrypt data using skipjack @re-fox
+  - encrypt data using camellia @re-fox
+  - hash data using tiger @re-fox
+  - encrypt data using blowfish @re-fox
+  - encrypt data using twofish @re-fox
+
+### Bug fixes
+
+  - linter: fix exception when examples is `None` @Ana06
+  - linter: fix suggested recommendations via templating @williballenthin
+  - render: fix exception when rendering counts @williballenthin
+  - render: fix render of negative offsets @williballenthin
+  - extractor: fix segmentation violation from vivisect @williballenthin
+  - main: fix crash when .viv cannot be saved #168 @secshoggoth @williballenthin
+  - main: fix shellcode .viv save path @williballenthin
+
+### Changes
+
+  - doc: explain how to bypass gatekeeper on macOS @psifertex
+  - doc: explain supported linux distributions @Ana06
+  - doc: explain submodule update with --init @psifertex
+  - main: improve program help output @mr-tz
+  - main: disable progress when run in quiet mode @mr-tz
+  - main: assert supported IDA versions @mr-tz
+  - extractor: better identify nested pointers to strings @williballenthin
+  - setup: specify vivisect download url @Ana06
+  - setup: pin vivisect version @williballenthin
+  - setup: bump vivisect dependency version @williballenthin
+  - setup: set Python project name to `flare-capa` @williballenthin
+  - ci: run tests and linter via Github Actions @Ana06
+  - hooks: run style checkers and hide stashed output @Ana06
+  - linter: ignore period in rule filename @williballenthin
+  - linter: warn on nursery rule with no changes needed @williballenthin
+
+### Raw diffs
+
+  - [capa v1.0.0...v1.1.0](https://github.com/fireeye/capa/compare/v1.0.0...v1.1.0)
+  - [capa-rules v1.0.0...v1.1.0](https://github.com/fireeye/capa-rules/compare/v1.0.0...v1.1.0)

LICENSE.txt

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright (C) 2020 FireEye, Inc.
+
+   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
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   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.

README.md

@@ -1,11 +1,15 @@
 ![capa](.github/logo.png)
 
-[![Number of rules](https://img.shields.io/badge/rules-261-blue.svg)](https://github.com/fireeye/capa-rules)
+[![CI status](https://github.com/fireeye/capa/workflows/CI/badge.svg)](https://github.com/fireeye/capa/actions?query=workflow%3ACI+event%3Apush+branch%3Amaster)
+[![Number of rules](https://img.shields.io/badge/rules-345-blue.svg)](https://github.com/fireeye/capa-rules)
+[![License](https://img.shields.io/badge/license-Apache--2.0-green.svg)](LICENSE.txt)
 
 capa detects capabilities in executable files.
 You run it against a PE file or shellcode and it tells you what it thinks the program can do.
 For example, it might suggest that the file is a backdoor, is capable of installing services, or relies on HTTP to communicate.
 
+Check out the overview in our first [capa blog post](https://www.fireeye.com/blog/threat-research/2020/07/capa-automatically-identify-malware-capabilities.html).
+
 ```
 $ capa.exe suspicious.exe
 
@@ -56,16 +60,18 @@ $ capa.exe suspicious.exe
 
 # download and usage
 
-Download stable releases of the standalone capa binaries [here](/releases). You can run the standalone binaries without installation.
+Download stable releases of the standalone capa binaries [here](https://github.com/fireeye/capa/releases). You can run the standalone binaries without installation. capa is a command line tool that should be run from the terminal.
 
+<!--
 Alternatively, you can fetch a nightly build of a standalone binary from one of the following links. These are built using the latest development branch.
 - Windows 64bit: TODO
 - Linux: TODO
 - OSX: TODO
+-->
 
 To use capa as a library or integrate with another tool, see [doc/installation.md](doc/installation.md) for further setup instructions.
 
-For more information about how to use capa, including running it as an IDA script/plugin see [doc/usage.md](doc/usage.md).
+For more information about how to use capa, see [doc/usage.md](doc/usage.md).
 
 # example
 
@@ -140,12 +146,11 @@ rule:
 The [github.com/fireeye/capa-rules](https://github.com/fireeye/capa-rules) repository contains hundreds of standard library rules that are distributed with capa.
 Please learn to write rules and contribute new entries as you find interesting techniques in malware.
 
-If you use IDA Pro, then you use can use the [IDA Pro plugin for capa](./capa/ida/ida_capa_explorer.py).
-This script adds new user interface elements to IDA, including an interactive tree view of rule matches and their locations within the current database.
-As you select the checkboxes, the plugin will highlight the addresses associated with the features.
-We use this plugin all the time to quickly jump to interesting parts of a program.
+If you use IDA Pro, then you use can use the [capa explorer IDA plugin](capa/ida/plugin/).
+capa explorer lets you quickly identify and navigate to interesting areas of a program and dissect capa rule matches at
+the assembly level.
 
-![capa + IDA Pro integration](.github/capa-ida.jpg)
+![capa + IDA Pro integration](doc/img/ida_plugin_intro.gif)
 
 # further information
 ## capa

capa/engine.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 sys
 import copy
@@ -14,12 +20,16 @@ class Statement(object):
      and to declare the interface method `evaluate`
     """
 
-    def __init__(self):
+    def __init__(self, description=None):
         super(Statement, self).__init__()
         self.name = self.__class__.__name__
+        self.description = description
 
     def __str__(self):
-        return "%s(%s)" % (self.name.lower(), ",".join(map(str, self.get_children())))
+        if self.description:
+            return "%s(%s = %s)" % (self.name.lower(), ",".join(map(str, self.get_children())), self.description)
+        else:
+            return "%s(%s)" % (self.name.lower(), ",".join(map(str, self.get_children())))
 
     def __repr__(self):
         return str(self)
@@ -98,9 +108,9 @@ class Result(object):
 class And(Statement):
     """match if all of the children evaluate to True."""
 
-    def __init__(self, *children):
-        super(And, self).__init__()
-        self.children = list(children)
+    def __init__(self, children, description=None):
+        super(And, self).__init__(description=description)
+        self.children = children
 
     def evaluate(self, ctx):
         results = [child.evaluate(ctx) for child in self.children]
@@ -111,9 +121,9 @@ class And(Statement):
 class Or(Statement):
     """match if any of the children evaluate to True."""
 
-    def __init__(self, *children):
-        super(Or, self).__init__()
-        self.children = list(children)
+    def __init__(self, children, description=None):
+        super(Or, self).__init__(description=description)
+        self.children = children
 
     def evaluate(self, ctx):
         results = [child.evaluate(ctx) for child in self.children]
@@ -124,8 +134,8 @@ class Or(Statement):
 class Not(Statement):
     """match only if the child evaluates to False."""
 
-    def __init__(self, child):
-        super(Not, self).__init__()
+    def __init__(self, child, description=None):
+        super(Not, self).__init__(description=description)
         self.child = child
 
     def evaluate(self, ctx):
@@ -137,10 +147,10 @@ class Not(Statement):
 class Some(Statement):
     """match if at least N of the children evaluate to True."""
 
-    def __init__(self, count, *children):
-        super(Some, self).__init__()
+    def __init__(self, count, children, description=None):
+        super(Some, self).__init__(description=description)
         self.count = count
-        self.children = list(children)
+        self.children = children
 
     def evaluate(self, ctx):
         results = [child.evaluate(ctx) for child in self.children]
@@ -155,8 +165,8 @@ class Some(Statement):
 class Range(Statement):
     """match if the child is contained in the ctx set with a count in the given range."""
 
-    def __init__(self, child, min=None, max=None):
-        super(Range, self).__init__()
+    def __init__(self, child, min=None, max=None, description=None):
+        super(Range, self).__init__(description=description)
         self.child = child
         self.min = min if min is not None else 0
         self.max = max if max is not None else (1 << 64 - 1)

capa/features/__init__.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 re
 import sys
@@ -10,6 +16,12 @@ import capa.engine
 logger = logging.getLogger(__name__)
 MAX_BYTES_FEATURE_SIZE = 0x100
 
+# identifiers for supported architectures names that tweak a feature
+# for example, offset/x32
+ARCH_X32 = "x32"
+ARCH_X64 = "x64"
+VALID_ARCH = (ARCH_X32, ARCH_X64)
+
 
 def bytes_to_str(b):
     if sys.version_info[0] >= 3:
@@ -24,21 +36,41 @@ def hex_string(h):
 
 
 class Feature(object):
-    def __init__(self, value, description=None):
+    def __init__(self, value, arch=None, description=None):
+        """
+        Args:
+          value (any): the value of the feature, such as the number or string.
+          arch (str): one of the VALID_ARCH values, or None.
+            When None, then the feature applies to any architecture.
+            Modifies the feature name from `feature` to `feature/arch`, like `offset/x32`.
+          description (str): a human-readable description that explains the feature value.
+        """
         super(Feature, self).__init__()
-        self.name = self.__class__.__name__.lower()
+
+        if arch is not None:
+            if arch not in VALID_ARCH:
+                raise ValueError("arch '%s' must be one of %s" % (arch, VALID_ARCH))
+            self.name = self.__class__.__name__.lower() + "/" + arch
+        else:
+            self.name = self.__class__.__name__.lower()
+
         self.value = value
+        self.arch = arch
         self.description = description
 
     def __hash__(self):
-        return hash((self.name, self.value))
+        return hash((self.name, self.value, self.arch))
 
     def __eq__(self, other):
-        return self.name == other.name and self.value == other.value
+        return self.name == other.name and self.value == other.value and self.arch == other.arch
 
-    # Used to overwrite the rendering of the feature value in `__str__` and the
-    # json output
     def get_value_str(self):
+        """
+        render the value of this feature, for use by `__str__` and friends.
+        subclasses should override to customize the rendering.
+
+        Returns: any
+        """
         return self.value
 
     def __str__(self):
@@ -56,36 +88,44 @@ class Feature(object):
     def evaluate(self, ctx):
         return capa.engine.Result(self in ctx, self, [], locations=ctx.get(self, []))
 
-    def serialize(self):
-        return self.__dict__
-
     def freeze_serialize(self):
-        return (self.__class__.__name__, [self.value])
+        if self.arch is not None:
+            return (self.__class__.__name__, [self.value, {"arch": self.arch}])
+        else:
+            return (self.__class__.__name__, [self.value])
 
     @classmethod
     def freeze_deserialize(cls, args):
-        return cls(*args)
+        # as you can see below in code,
+        # if the last argument is a dictionary,
+        # consider it to be kwargs passed to the feature constructor.
+        if len(args) == 1:
+            return cls(*args)
+        elif isinstance(args[-1], dict):
+            kwargs = args[-1]
+            args = args[:-1]
+            return cls(*args, **kwargs)
 
 
 class MatchedRule(Feature):
     def __init__(self, value, description=None):
-        super(MatchedRule, self).__init__(value, description)
+        super(MatchedRule, self).__init__(value, description=description)
         self.name = "match"
 
 
 class Characteristic(Feature):
     def __init__(self, value, description=None):
-        super(Characteristic, self).__init__(value, description)
+        super(Characteristic, self).__init__(value, description=description)
 
 
 class String(Feature):
     def __init__(self, value, description=None):
-        super(String, self).__init__(value, description)
+        super(String, self).__init__(value, description=description)
 
 
 class Regex(String):
     def __init__(self, value, description=None):
-        super(Regex, self).__init__(value, description)
+        super(Regex, self).__init__(value, description=description)
         pat = self.value[len("/") : -len("/")]
         flags = re.DOTALL
         if value.endswith("/i"):
@@ -99,7 +139,6 @@ class Regex(String):
             raise ValueError(
                 "invalid regular expression: %s it should use Python syntax, try it at https://pythex.org" % value
             )
-        self.match = None
 
     def evaluate(self, ctx):
         for feature, locations in ctx.items():
@@ -111,25 +150,53 @@ class Regex(String):
             # using this mode cleans is more convenient for rule authors,
             # so that they don't have to prefix/suffix their terms like: /.*foo.*/.
             if self.re.search(feature.value):
-                self.match = feature.value
-                return capa.engine.Result(True, self, [], locations=locations)
+                # unlike other features, we cannot return put a reference to `self` directly in a `Result`.
+                # this is because `self` may match on many strings, so we can't stuff the matched value into it.
+                # instead, return a new instance that has a reference to both the regex and the matched value.
+                # see #262.
+                return capa.engine.Result(True, _MatchedRegex(self, feature.value), [], locations=locations)
 
-        return capa.engine.Result(False, self, [])
+        return capa.engine.Result(False, _MatchedRegex(self, None), [])
+
+    def __str__(self):
+        return "regex(string =~ %s)" % self.value
+
+
+class _MatchedRegex(Regex):
+    """
+    this represents a specific instance of a regular expression feature match.
+    treat it the same as a `Regex` except it has the `match` field that contains the complete string that matched.
+
+    note: this type should only ever be constructed by `Regex.evaluate()`. it is not part of the public API.
+    """
+
+    def __init__(self, regex, match):
+        """
+        args:
+          regex (Regex): the regex feature that matches
+          match (string|None): the matching string or None if it doesn't match
+        """
+        super(_MatchedRegex, self).__init__(regex.value, description=regex.description)
+        # we want this to collide with the name of `Regex` above,
+        # so that it works nicely with the renderers.
+        self.name = "regex"
+        # this may be None if the regex doesn't match
+        self.match = match
 
     def __str__(self):
         return 'regex(string =~ %s, matched = "%s")' % (self.value, self.match)
 
 
 class StringFactory(object):
-    def __new__(self, value, description):
+    def __new__(self, value, description=None):
         if value.startswith("/") and (value.endswith("/") or value.endswith("/i")):
-            return Regex(value, description)
-        return String(value, description)
+            return Regex(value, description=description)
+        return String(value, description=description)
 
 
 class Bytes(Feature):
     def __init__(self, value, description=None):
-        super(Bytes, self).__init__(value, description)
+        super(Bytes, self).__init__(value, description=description)
 
     def evaluate(self, ctx):
         for feature, locations in ctx.items():

capa/features/basicblock.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 capa.features import Feature
 

capa/features/extractors/__init__.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 abc
 
@@ -190,7 +196,7 @@ class NullFeatureExtractor(FeatureExtractor):
             'functions': {
                 0x401000: {
                     'features': [
-                        (0x401000, capa.features.Characteristic('switch')),
+                        (0x401000, capa.features.Characteristic('nzxor')),
                     ],
                     'basic blocks': {
                         0x401000: {

capa/features/extractors/helpers.py

@@ -1,8 +1,15 @@
 # Copyright (C) 2020 FireEye, 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 sys
 import builtins
 
+from capa.features.file import Import
 from capa.features.insn import API
 
 MIN_STACKSTRING_LEN = 8
@@ -15,25 +22,32 @@ def xor_static(data, i):
         return "".join(chr(ord(c) ^ i) for c in data)
 
 
-def is_aw_function(function_name):
+def is_aw_function(symbol):
     """
     is the given function name an A/W function?
     these are variants of functions that, on Windows, accept either a narrow or wide string.
     """
-    if len(function_name) < 2:
+    if len(symbol) < 2:
         return False
 
     # last character should be 'A' or 'W'
-    if function_name[-1] not in ("A", "W"):
+    if symbol[-1] not in ("A", "W"):
         return False
 
     # second to last character should be lowercase letter
-    return "a" <= function_name[-2] <= "z" or "0" <= function_name[-2] <= "9"
+    return "a" <= symbol[-2] <= "z" or "0" <= symbol[-2] <= "9"
 
 
-def generate_api_features(apiname, va):
+def is_ordinal(symbol):
     """
-    for a given function name and address, generate API names.
+    is the given symbol an ordinal that is prefixed by "#"?
+    """
+    return symbol[0] == "#"
+
+
+def generate_symbols(dll, symbol):
+    """
+    for a given dll and symbol name, generate variants.
     we over-generate features to make matching easier.
     these include:
       - kernel32.CreateFileA
@@ -41,23 +55,36 @@ def generate_api_features(apiname, va):
       - CreateFileA
       - CreateFile
     """
-    # (kernel32.CreateFileA, 0x401000)
-    yield API(apiname), va
+    # kernel32.CreateFileA
+    yield "%s.%s" % (dll, symbol)
 
-    if is_aw_function(apiname):
-        # (kernel32.CreateFile, 0x401000)
-        yield API(apiname[:-1]), va
+    if not is_ordinal(symbol):
+        # CreateFileA
+        yield symbol
 
-    if "." in apiname:
-        modname, impname = apiname.split(".")
-        # strip modname to support importname-only matching
-        # (CreateFileA, 0x401000)
-        yield API(impname), va
+    if is_aw_function(symbol):
+        # kernel32.CreateFile
+        yield "%s.%s" % (dll, symbol[:-1])
 
-        if is_aw_function(impname):
-            # (CreateFile, 0x401000)
-            yield API(impname[:-1]), va
+        if not is_ordinal(symbol):
+            # CreateFile
+            yield symbol[:-1]
 
 
 def all_zeros(bytez):
     return all(b == 0 for b in builtins.bytes(bytez))
+
+
+def twos_complement(val, bits):
+    """
+    compute the 2's complement of int value val
+
+    from: https://stackoverflow.com/a/9147327/87207
+    """
+    # if sign bit is set e.g., 8bit: 128-255
+    if (val & (1 << (bits - 1))) != 0:
+        # compute negative value
+        return val - (1 << bits)
+    else:
+        # return positive value as is
+        return val

capa/features/extractors/ida/__init__.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 sys
 import types
@@ -49,16 +55,27 @@ class IdaFeatureExtractor(FeatureExtractor):
     def get_functions(self):
         import capa.features.extractors.ida.helpers as ida_helpers
 
+        # data structure shared across functions yielded here.
+        # useful for caching analysis relevant across a single workspace.
+        ctx = {}
+
         # ignore library functions and thunk functions as identified by IDA
         for f in ida_helpers.get_functions(skip_thunks=True, skip_libs=True):
+            setattr(f, "ctx", ctx)
             yield add_ea_int_cast(f)
 
+    @staticmethod
+    def get_function(ea):
+        f = idaapi.get_func(ea)
+        setattr(f, "ctx", {})
+        return add_ea_int_cast(f)
+
     def extract_function_features(self, f):
         for (feature, ea) in capa.features.extractors.ida.function.extract_features(f):
             yield feature, ea
 
     def get_basic_blocks(self, f):
-        for bb in idaapi.FlowChart(f, flags=idaapi.FC_PREDS):
+        for bb in capa.features.extractors.ida.helpers.get_function_blocks(f):
             yield add_ea_int_cast(bb)
 
     def extract_basic_block_features(self, f, bb):

capa/features/extractors/ida/basicblock.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 sys
 import string
@@ -14,10 +20,10 @@ from capa.features.extractors.helpers import MIN_STACKSTRING_LEN
 
 
 def get_printable_len(op):
-    """ Return string length if all operand bytes are ascii or utf16-le printable
+    """Return string length if all operand bytes are ascii or utf16-le printable
 
-        args:
-            op (IDA op_t)
+    args:
+        op (IDA op_t)
     """
     op_val = capa.features.extractors.ida.helpers.mask_op_val(op)
 
@@ -56,10 +62,10 @@ def get_printable_len(op):
 
 
 def is_mov_imm_to_stack(insn):
-    """ verify instruction moves immediate onto stack
+    """verify instruction moves immediate onto stack
 
-        args:
-            insn (IDA insn_t)
+    args:
+        insn (IDA insn_t)
     """
     if insn.Op2.type != idaapi.o_imm:
         return False
@@ -74,13 +80,13 @@ def is_mov_imm_to_stack(insn):
 
 
 def bb_contains_stackstring(f, bb):
-    """ check basic block for stackstring indicators
+    """check basic block for stackstring indicators
 
-        true if basic block contains enough moves of constant bytes to the stack
+    true if basic block contains enough moves of constant bytes to the stack
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
     """
     count = 0
     for insn in capa.features.extractors.ida.helpers.get_instructions_in_range(bb.start_ea, bb.end_ea):
@@ -92,33 +98,33 @@ def bb_contains_stackstring(f, bb):
 
 
 def extract_bb_stackstring(f, bb):
-    """ extract stackstring indicators from basic block
+    """extract stackstring indicators from basic block
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
     """
     if bb_contains_stackstring(f, bb):
         yield Characteristic("stack string"), bb.start_ea
 
 
 def extract_bb_tight_loop(f, bb):
-    """ extract tight loop indicators from a basic block
+    """extract tight loop indicators from a basic block
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
     """
     if capa.features.extractors.ida.helpers.is_basic_block_tight_loop(bb):
         yield Characteristic("tight loop"), bb.start_ea
 
 
 def extract_features(f, bb):
-    """ extract basic block features
+    """extract basic block features
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
     """
     for bb_handler in BASIC_BLOCK_HANDLERS:
         for (feature, ea) in bb_handler(f, bb):

capa/features/extractors/ida/file.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 struct
 
@@ -14,13 +20,13 @@ from capa.features.file import Export, Import, Section
 
 
 def check_segment_for_pe(seg):
-    """ check segment for embedded PE
+    """check segment for embedded PE
 
-        adapted for IDA from:
-        https://github.com/vivisect/vivisect/blob/7be4037b1cecc4551b397f840405a1fc606f9b53/PE/carve.py#L19
+    adapted for IDA from:
+    https://github.com/vivisect/vivisect/blob/7be4037b1cecc4551b397f840405a1fc606f9b53/PE/carve.py#L19
 
-        args:
-            seg (IDA segment_t)
+    args:
+        seg (IDA segment_t)
     """
     seg_max = seg.end_ea
     mz_xor = [
@@ -61,11 +67,11 @@ def check_segment_for_pe(seg):
 
 
 def extract_file_embedded_pe():
-    """ extract embedded PE features
+    """extract embedded PE features
 
-        IDA must load resource sections for this to be complete
-            - '-R' from console
-            - Check 'Load resource sections' when opening binary in IDA manually
+    IDA must load resource sections for this to be complete
+        - '-R' from console
+        - Check 'Load resource sections' when opening binary in IDA manually
     """
     for seg in capa.features.extractors.ida.helpers.get_segments(skip_header_segments=True):
         for (ea, _) in check_segment_for_pe(seg):
@@ -79,41 +85,47 @@ def extract_file_export_names():
 
 
 def extract_file_import_names():
-    """ extract function imports
+    """extract function imports
 
-        1. imports by ordinal:
-         - modulename.#ordinal
+    1. imports by ordinal:
+     - modulename.#ordinal
 
-        2. imports by name, results in two features to support importname-only
-           matching:
-         - modulename.importname
-         - importname
+    2. imports by name, results in two features to support importname-only
+       matching:
+     - modulename.importname
+     - importname
     """
     for (ea, info) in capa.features.extractors.ida.helpers.get_file_imports().items():
         if info[1]:
-            yield Import("%s.%s" % (info[0], info[1])), ea
-            yield Import(info[1]), ea
-        if info[2]:
-            yield Import("%s.#%s" % (info[0], str(info[2]))), ea
+            dll = info[0]
+            symbol = info[1]
+        elif info[2]:
+            dll = info[0]
+            symbol = "#%d" % (info[2])
+        else:
+            continue
+
+        for name in capa.features.extractors.helpers.generate_symbols(dll, symbol):
+            yield Import(name), ea
 
 
 def extract_file_section_names():
-    """ extract section names
+    """extract section names
 
-        IDA must load resource sections for this to be complete
-            - '-R' from console
-            - Check 'Load resource sections' when opening binary in IDA manually
+    IDA must load resource sections for this to be complete
+        - '-R' from console
+        - Check 'Load resource sections' when opening binary in IDA manually
     """
     for seg in capa.features.extractors.ida.helpers.get_segments(skip_header_segments=True):
         yield Section(idaapi.get_segm_name(seg)), seg.start_ea
 
 
 def extract_file_strings():
-    """ extract ASCII and UTF-16 LE strings
+    """extract ASCII and UTF-16 LE strings
 
-        IDA must load resource sections for this to be complete
-            - '-R' from console
-            - Check 'Load resource sections' when opening binary in IDA manually
+    IDA must load resource sections for this to be complete
+        - '-R' from console
+        - Check 'Load resource sections' when opening binary in IDA manually
     """
     for seg in capa.features.extractors.ida.helpers.get_segments():
         seg_buff = capa.features.extractors.ida.helpers.get_segment_buffer(seg)

capa/features/extractors/ida/function.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 idaapi
 import idautils
@@ -8,31 +14,21 @@ from capa.features import Characteristic
 from capa.features.extractors import loops
 
 
-def extract_function_switch(f):
-    """ extract switch indicators from a function
-
-        arg:
-            f (IDA func_t)
-    """
-    if capa.features.extractors.ida.helpers.is_function_switch_statement(f):
-        yield Characteristic("switch"), f.start_ea
-
-
 def extract_function_calls_to(f):
-    """ extract callers to a function
+    """extract callers to a function
 
-        args:
-            f (IDA func_t)
+    args:
+        f (IDA func_t)
     """
     for ea in idautils.CodeRefsTo(f.start_ea, True):
         yield Characteristic("calls to"), ea
 
 
 def extract_function_loop(f):
-    """ extract loop indicators from a function
+    """extract loop indicators from a function
 
-        args:
-            f (IDA func_t)
+    args:
+        f (IDA func_t)
     """
     edges = []
 
@@ -46,27 +42,27 @@ def extract_function_loop(f):
 
 
 def extract_recursive_call(f):
-    """ extract recursive function call
+    """extract recursive function call
 
-        args:
-            f (IDA func_t)
+    args:
+        f (IDA func_t)
     """
     if capa.features.extractors.ida.helpers.is_function_recursive(f):
         yield Characteristic("recursive call"), f.start_ea
 
 
 def extract_features(f):
-    """ extract function features
+    """extract function features
 
-        arg:
-            f (IDA func_t)
+    arg:
+        f (IDA func_t)
     """
     for func_handler in FUNCTION_HANDLERS:
         for (feature, ea) in func_handler(f):
             yield feature, ea
 
 
-FUNCTION_HANDLERS = (extract_function_calls_to, extract_function_switch, extract_function_loop, extract_recursive_call)
+FUNCTION_HANDLERS = (extract_function_calls_to, extract_function_loop, extract_recursive_call)
 
 
 def main():

capa/features/extractors/ida/helpers.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 sys
 import string
@@ -6,15 +12,16 @@ import string
 import idc
 import idaapi
 import idautils
+import ida_bytes
 
 
 def find_byte_sequence(start, end, seq):
-    """ find byte sequence
+    """find byte sequence
 
-        args:
-            start: min virtual address
-            end: max virtual address
-            seq: bytes to search e.g. b'\x01\x03'
+    args:
+        start: min virtual address
+        end: max virtual address
+        seq: bytes to search e.g. b'\x01\x03'
     """
     if sys.version_info[0] >= 3:
         return idaapi.find_binary(start, end, " ".join(["%02x" % b for b in seq]), 0, idaapi.SEARCH_DOWN)
@@ -23,14 +30,14 @@ def find_byte_sequence(start, end, seq):
 
 
 def get_functions(start=None, end=None, skip_thunks=False, skip_libs=False):
-    """ get functions, range optional
+    """get functions, range optional
 
-        args:
-            start: min virtual address
-            end: max virtual address
+    args:
+        start: min virtual address
+        end: max virtual address
 
-        ret:
-            yield func_t*
+    ret:
+        yield func_t*
     """
     for ea in idautils.Functions(start=start, end=end):
         f = idaapi.get_func(ea)
@@ -39,10 +46,10 @@ def get_functions(start=None, end=None, skip_thunks=False, skip_libs=False):
 
 
 def get_segments(skip_header_segments=False):
-    """ get list of segments (sections) in the binary image
+    """get list of segments (sections) in the binary image
 
-        args:
-            skip_header_segments: IDA may load header segments - skip if set
+    args:
+        skip_header_segments: IDA may load header segments - skip if set
     """
     for n in range(idaapi.get_segm_qty()):
         seg = idaapi.getnseg(n)
@@ -51,9 +58,9 @@ def get_segments(skip_header_segments=False):
 
 
 def get_segment_buffer(seg):
-    """ return bytes stored in a given segment
+    """return bytes stored in a given segment
 
-        decrease buffer size until IDA is able to read bytes from the segment
+    decrease buffer size until IDA is able to read bytes from the segment
     """
     buff = b""
     sz = seg.end_ea - seg.start_ea
@@ -91,13 +98,13 @@ def get_file_imports():
 
 
 def get_instructions_in_range(start, end):
-    """ yield instructions in range
+    """yield instructions in range
 
-        args:
-            start: virtual address (inclusive)
-            end: virtual address (exclusive)
-        yield:
-            (insn_t*)
+    args:
+        start: virtual address (inclusive)
+        end: virtual address (exclusive)
+    yield:
+        (insn_t*)
     """
     for head in idautils.Heads(start, end):
         insn = idautils.DecodeInstruction(head)
@@ -177,10 +184,10 @@ def find_string_at(ea, min=4):
 
 
 def get_op_phrase_info(op):
-    """ parse phrase features from operand
+    """parse phrase features from operand
 
-        Pretty much dup of sark's implementation:
-            https://github.com/tmr232/Sark/blob/master/sark/code/instruction.py#L28-L73
+    Pretty much dup of sark's implementation:
+        https://github.com/tmr232/Sark/blob/master/sark/code/instruction.py#L28-L73
     """
     if op.type not in (idaapi.o_phrase, idaapi.o_displ):
         return {}
@@ -223,6 +230,12 @@ def is_op_read(insn, op):
     return idaapi.has_cf_use(insn.get_canon_feature(), op.n)
 
 
+def is_op_offset(insn, op):
+    """ Check is an operand has been marked as an offset (by auto-analysis or manually) """
+    flags = idaapi.get_flags(insn.ea)
+    return ida_bytes.is_off(flags, op.n)
+
+
 def is_sp_modified(insn):
     """ determine if instruction modifies SP, ESP, RSP """
     for op in get_insn_ops(insn, target_ops=(idaapi.o_reg,)):
@@ -263,15 +276,15 @@ def is_op_stack_var(ea, index):
 
 
 def mask_op_val(op):
-    """ mask value by data type
+    """mask value by data type
 
-        necessary due to a bug in AMD64
+    necessary due to a bug in AMD64
 
-        Example:
-            .rsrc:0054C12C mov [ebp+var_4], 0FFFFFFFFh
+    Example:
+        .rsrc:0054C12C mov [ebp+var_4], 0FFFFFFFFh
 
-            insn.Op2.dtype == idaapi.dt_dword
-            insn.Op2.value == 0xffffffffffffffff
+        insn.Op2.dtype == idaapi.dt_dword
+        insn.Op2.value == 0xffffffffffffffff
     """
     masks = {
         idaapi.dt_byte: 0xFF,
@@ -283,10 +296,10 @@ def mask_op_val(op):
 
 
 def is_function_recursive(f):
-    """ check if function is recursive
+    """check if function is recursive
 
-        args:
-            f (IDA func_t)
+    args:
+        f (IDA func_t)
     """
     for ref in idautils.CodeRefsTo(f.start_ea, True):
         if f.contains(ref):
@@ -294,30 +307,14 @@ def is_function_recursive(f):
     return False
 
 
-def is_function_switch_statement(f):
-    """ check a function for switch statement indicators
-
-        adapted from:
-        https://reverseengineering.stackexchange.com/questions/17548/calc-switch-cases-in-idapython-cant-iterate-over-results?rq=1
-
-        arg:
-            f (IDA func_t)
-    """
-    for (start, end) in idautils.Chunks(f.start_ea):
-        for head in idautils.Heads(start, end):
-            if idaapi.get_switch_info(head):
-                return True
-    return False
-
-
 def is_basic_block_tight_loop(bb):
-    """ check basic block loops to self
+    """check basic block loops to self
 
-        true if last instruction in basic block branches to basic block start
+    true if last instruction in basic block branches to basic block start
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
     """
     bb_end = idc.prev_head(bb.end_ea)
     if bb.start_ea < bb_end:
@@ -325,3 +322,47 @@ def is_basic_block_tight_loop(bb):
             if ref == bb.start_ea:
                 return True
     return False
+
+
+def find_data_reference_from_insn(insn, max_depth=10):
+    """ search for data reference from instruction, return address of instruction if no reference exists """
+    depth = 0
+    ea = insn.ea
+
+    while True:
+        data_refs = list(idautils.DataRefsFrom(ea))
+
+        if len(data_refs) != 1:
+            # break if no refs or more than one ref (assume nested pointers only have one data reference)
+            break
+
+        if ea == data_refs[0]:
+            # break if circular reference
+            break
+
+        depth += 1
+        if depth > max_depth:
+            # break if max depth
+            break
+
+        ea = data_refs[0]
+
+    return ea
+
+
+def get_function_blocks(f):
+    """yield basic blocks contained in specified function
+
+    args:
+        f (IDA func_t)
+    yield:
+        block (IDA BasicBlock)
+    """
+    # leverage idaapi.FC_NOEXT flag to ignore useless external blocks referenced by the function
+    for block in idaapi.FlowChart(f, flags=(idaapi.FC_PREDS | idaapi.FC_NOEXT)):
+        yield block
+
+
+def is_basic_block_return(bb):
+    """ check if basic block is return block """
+    return bb.type == idaapi.fcb_ret

capa/features/extractors/ida/insn.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 idc
 import idaapi
@@ -6,27 +12,45 @@ import idautils
 
 import capa.features.extractors.helpers
 import capa.features.extractors.ida.helpers
-from capa.features import MAX_BYTES_FEATURE_SIZE, Bytes, String, Characteristic
-from capa.features.insn import Number, Offset, Mnemonic
+from capa.features import ARCH_X32, ARCH_X64, MAX_BYTES_FEATURE_SIZE, Bytes, String, Characteristic
+from capa.features.insn import API, Number, Offset, Mnemonic
 
-_file_imports_cache = None
+# security cookie checks may perform non-zeroing XORs, these are expected within a certain
+# byte range within the first and returning basic blocks, this helps to reduce FP features
+SECURITY_COOKIE_BYTES_DELTA = 0x40
 
 
-def get_imports():
-    """ """
-    global _file_imports_cache
-    if _file_imports_cache is None:
-        _file_imports_cache = capa.features.extractors.ida.helpers.get_file_imports()
-    return _file_imports_cache
+def get_arch(ctx):
+    """
+    fetch the ARCH_* constant for the currently open workspace.
+
+    via Tamir Bahar/@tmr232
+    https://reverseengineering.stackexchange.com/a/11398/17194
+    """
+    if "arch" not in ctx:
+        info = idaapi.get_inf_structure()
+        if info.is_64bit():
+            ctx["arch"] = ARCH_X64
+        elif info.is_32bit():
+            ctx["arch"] = ARCH_X32
+        else:
+            raise ValueError("unexpected architecture")
+    return ctx["arch"]
 
 
-def check_for_api_call(insn):
+def get_imports(ctx):
+    if "imports_cache" not in ctx:
+        ctx["imports_cache"] = capa.features.extractors.ida.helpers.get_file_imports()
+    return ctx["imports_cache"]
+
+
+def check_for_api_call(ctx, insn):
     """ check instruction for API call """
     if not idaapi.is_call_insn(insn):
         return
 
     for ref in idautils.CodeRefsFrom(insn.ea, False):
-        info = get_imports().get(ref, ())
+        info = get_imports(ctx).get(ref, ())
         if info:
             yield "%s.%s" % (info[0], info[1])
         else:
@@ -36,37 +60,38 @@ def check_for_api_call(insn):
             if f and (f.flags & idaapi.FUNC_THUNK):
                 for thunk_ref in idautils.DataRefsFrom(ref):
                     # TODO: always data ref for thunk??
-                    info = get_imports().get(thunk_ref, ())
+                    info = get_imports(ctx).get(thunk_ref, ())
                     if info:
                         yield "%s.%s" % (info[0], info[1])
 
 
 def extract_insn_api_features(f, bb, insn):
-    """ parse instruction API features
+    """parse instruction API features
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
 
-        example:
-            call dword [0x00473038]
+    example:
+        call dword [0x00473038]
     """
-    for api in check_for_api_call(insn):
-        for (feature, ea) in capa.features.extractors.helpers.generate_api_features(api, insn.ea):
-            yield feature, ea
+    for api in check_for_api_call(f.ctx, insn):
+        dll, _, symbol = api.rpartition(".")
+        for name in capa.features.extractors.helpers.generate_symbols(dll, symbol):
+            yield API(name), insn.ea
 
 
 def extract_insn_number_features(f, bb, insn):
-    """ parse instruction number features
+    """parse instruction number features
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
 
-        example:
-            push    3136B0h         ; dwControlCode
+    example:
+        push    3136B0h         ; dwControlCode
     """
     if idaapi.is_ret_insn(insn):
         # skip things like:
@@ -78,81 +103,93 @@ def extract_insn_number_features(f, bb, insn):
         #   .text:00401145 add esp, 0Ch
         return
 
-    for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, target_ops=(idaapi.o_imm,)):
-        const = capa.features.extractors.ida.helpers.mask_op_val(op)
-        if not idaapi.is_mapped(const):
-            yield Number(const), insn.ea
+    for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, target_ops=(idaapi.o_imm, idaapi.o_mem)):
+        # skip things like:
+        #   .text:00401100 shr eax, offset loc_C
+        if capa.features.extractors.ida.helpers.is_op_offset(insn, op):
+            continue
+
+        if op.type == idaapi.o_imm:
+            const = capa.features.extractors.ida.helpers.mask_op_val(op)
+        else:
+            const = op.addr
+
+        yield Number(const), insn.ea
+        yield Number(const, arch=get_arch(f.ctx)), insn.ea
 
 
 def extract_insn_bytes_features(f, bb, insn):
-    """ parse referenced byte sequences
+    """parse referenced byte sequences
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
 
-        example:
-            push    offset iid_004118d4_IShellLinkA ; riid
+    example:
+        push    offset iid_004118d4_IShellLinkA ; riid
     """
-    if idaapi.is_call_insn(insn):
-        # ignore call instructions
-        return
-
-    for ref in idautils.DataRefsFrom(insn.ea):
+    ref = capa.features.extractors.ida.helpers.find_data_reference_from_insn(insn)
+    if ref != insn.ea:
         extracted_bytes = capa.features.extractors.ida.helpers.read_bytes_at(ref, MAX_BYTES_FEATURE_SIZE)
         if extracted_bytes and not capa.features.extractors.helpers.all_zeros(extracted_bytes):
             yield Bytes(extracted_bytes), insn.ea
 
 
 def extract_insn_string_features(f, bb, insn):
-    """ parse instruction string features
+    """parse instruction string features
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
 
-        example:
-            push offset aAcr     ; "ACR  > "
+    example:
+        push offset aAcr     ; "ACR  > "
     """
-    for ref in idautils.DataRefsFrom(insn.ea):
+    ref = capa.features.extractors.ida.helpers.find_data_reference_from_insn(insn)
+    if ref != insn.ea:
         found = capa.features.extractors.ida.helpers.find_string_at(ref)
         if found:
             yield String(found), insn.ea
 
 
 def extract_insn_offset_features(f, bb, insn):
-    """ parse instruction structure offset features
+    """parse instruction structure offset features
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
 
-        example:
-            .text:0040112F cmp [esi+4], ebx
+    example:
+        .text:0040112F cmp [esi+4], ebx
     """
     for op in capa.features.extractors.ida.helpers.get_insn_ops(insn, target_ops=(idaapi.o_phrase, idaapi.o_displ)):
         if capa.features.extractors.ida.helpers.is_op_stack_var(insn.ea, op.n):
             continue
         p_info = capa.features.extractors.ida.helpers.get_op_phrase_info(op)
         op_off = p_info.get("offset", 0)
-        if 0 == op_off:
-            continue
         if idaapi.is_mapped(op_off):
             # Ignore:
             #   mov esi, dword_1005B148[esi]
             continue
+
+        # I believe that IDA encodes all offsets as two's complement in a u32.
+        # a 64-bit displacement isn't a thing, see:
+        # https://stackoverflow.com/questions/31853189/x86-64-assembly-why-displacement-not-64-bits
+        op_off = capa.features.extractors.helpers.twos_complement(op_off, 32)
+
         yield Offset(op_off), insn.ea
+        yield Offset(op_off, arch=get_arch(f.ctx)), insn.ea
 
 
 def contains_stack_cookie_keywords(s):
-    """ check if string contains stack cookie keywords
+    """check if string contains stack cookie keywords
 
-        Examples:
-            xor     ecx, ebp ; StackCookie
-            mov     eax, ___security_cookie
+    Examples:
+        xor     ecx, ebp ; StackCookie
+        mov     eax, ___security_cookie
     """
     if not s:
         return False
@@ -163,30 +200,30 @@ def contains_stack_cookie_keywords(s):
 
 
 def bb_stack_cookie_registers(bb):
-    """ scan basic block for stack cookie operations
+    """scan basic block for stack cookie operations
 
-        yield registers ids that may have been used for stack cookie operations
+    yield registers ids that may have been used for stack cookie operations
 
-        assume instruction that sets stack cookie and nzxor exist in same block
-        and stack cookie register is not modified prior to nzxor
+    assume instruction that sets stack cookie and nzxor exist in same block
+    and stack cookie register is not modified prior to nzxor
 
-        Example:
-            .text:004062DA mov     eax, ___security_cookie <-- stack cookie
-            .text:004062DF mov     ecx, eax
-            .text:004062E1 mov     ebx, [esi]
-            .text:004062E3 and     ecx, 1Fh
-            .text:004062E6 mov     edi, [esi+4]
-            .text:004062E9 xor     ebx, eax
-            .text:004062EB mov     esi, [esi+8]
-            .text:004062EE xor     edi, eax <-- ignore
-            .text:004062F0 xor     esi, eax <-- ignore
-            .text:004062F2 ror     edi, cl
-            .text:004062F4 ror     esi, cl
-            .text:004062F6 ror     ebx, cl
-            .text:004062F8 cmp     edi, esi
-            .text:004062FA jnz     loc_40639D
+    Example:
+        .text:004062DA mov     eax, ___security_cookie <-- stack cookie
+        .text:004062DF mov     ecx, eax
+        .text:004062E1 mov     ebx, [esi]
+        .text:004062E3 and     ecx, 1Fh
+        .text:004062E6 mov     edi, [esi+4]
+        .text:004062E9 xor     ebx, eax
+        .text:004062EB mov     esi, [esi+8]
+        .text:004062EE xor     edi, eax <-- ignore
+        .text:004062F0 xor     esi, eax <-- ignore
+        .text:004062F2 ror     edi, cl
+        .text:004062F4 ror     esi, cl
+        .text:004062F6 ror     ebx, cl
+        .text:004062F8 cmp     edi, esi
+        .text:004062FA jnz     loc_40639D
 
-        TODO: this is expensive, but necessary?...
+    TODO: this is expensive, but necessary?...
     """
     for insn in capa.features.extractors.ida.helpers.get_instructions_in_range(bb.start_ea, bb.end_ea):
         if contains_stack_cookie_keywords(idc.GetDisasm(insn.ea)):
@@ -196,12 +233,37 @@ def bb_stack_cookie_registers(bb):
                     yield op.reg
 
 
+def is_nzxor_stack_cookie_delta(f, bb, insn):
+    """ check if nzxor exists within stack cookie delta """
+    # security cookie check should use SP or BP
+    if not capa.features.extractors.ida.helpers.is_frame_register(insn.Op2.reg):
+        return False
+
+    f_bbs = tuple(capa.features.extractors.ida.helpers.get_function_blocks(f))
+
+    # expect security cookie init in first basic block within first bytes (instructions)
+    if capa.features.extractors.ida.helpers.is_basic_block_equal(bb, f_bbs[0]) and insn.ea < (
+        bb.start_ea + SECURITY_COOKIE_BYTES_DELTA
+    ):
+        return True
+
+    # ... or within last bytes (instructions) before a return
+    if capa.features.extractors.ida.helpers.is_basic_block_return(bb) and insn.ea > (
+        bb.start_ea + capa.features.extractors.ida.helpers.basic_block_size(bb) - SECURITY_COOKIE_BYTES_DELTA
+    ):
+        return True
+
+    return False
+
+
 def is_nzxor_stack_cookie(f, bb, insn):
     """ check if nzxor is related to stack cookie """
     if contains_stack_cookie_keywords(idaapi.get_cmt(insn.ea, False)):
         # Example:
         #   xor     ecx, ebp        ; StackCookie
         return True
+    if is_nzxor_stack_cookie_delta(f, bb, insn):
+        return True
     stack_cookie_regs = tuple(bb_stack_cookie_registers(bb))
     if any(op_reg in stack_cookie_regs for op_reg in (insn.Op1.reg, insn.Op2.reg)):
         # Example:
@@ -212,14 +274,14 @@ def is_nzxor_stack_cookie(f, bb, insn):
 
 
 def extract_insn_nzxor_characteristic_features(f, bb, insn):
-    """ parse instruction non-zeroing XOR instruction
+    """parse instruction non-zeroing XOR instruction
 
-        ignore expected non-zeroing XORs, e.g. security cookies
+    ignore expected non-zeroing XORs, e.g. security cookies
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
     """
     if insn.itype != idaapi.NN_xor:
         return
@@ -231,23 +293,23 @@ def extract_insn_nzxor_characteristic_features(f, bb, insn):
 
 
 def extract_insn_mnemonic_features(f, bb, insn):
-    """ parse instruction mnemonic features
+    """parse instruction mnemonic features
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
     """
     yield Mnemonic(insn.get_canon_mnem()), insn.ea
 
 
 def extract_insn_peb_access_characteristic_features(f, bb, insn):
-    """ parse instruction peb access
+    """parse instruction peb access
 
-        fs:[0x30] on x86, gs:[0x60] on x64
+    fs:[0x30] on x86, gs:[0x60] on x64
 
-        TODO:
-            IDA should be able to do this..
+    TODO:
+        IDA should be able to do this..
     """
     if insn.itype not in (idaapi.NN_push, idaapi.NN_mov):
         return
@@ -264,10 +326,10 @@ def extract_insn_peb_access_characteristic_features(f, bb, insn):
 
 
 def extract_insn_segment_access_features(f, bb, insn):
-    """ parse instruction fs or gs access
+    """parse instruction fs or gs access
 
-        TODO:
-            IDA should be able to do this...
+    TODO:
+        IDA should be able to do this...
     """
     if all(map(lambda op: op.type != idaapi.o_mem, insn.ops)):
         # try to optimize for only memory references
@@ -285,15 +347,15 @@ def extract_insn_segment_access_features(f, bb, insn):
 
 
 def extract_insn_cross_section_cflow(f, bb, insn):
-    """ inspect the instruction for a CALL or JMP that crosses section boundaries
+    """inspect the instruction for a CALL or JMP that crosses section boundaries
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
     """
     for ref in idautils.CodeRefsFrom(insn.ea, False):
-        if ref in get_imports().keys():
+        if ref in get_imports(f.ctx).keys():
             # ignore API calls
             continue
         if not idaapi.getseg(ref):
@@ -305,14 +367,14 @@ def extract_insn_cross_section_cflow(f, bb, insn):
 
 
 def extract_function_calls_from(f, bb, insn):
-    """ extract functions calls from features
+    """extract functions calls from features
 
-        most relevant at the function scope, however, its most efficient to extract at the instruction scope
+    most relevant at the function scope, however, its most efficient to extract at the instruction scope
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
     """
     if idaapi.is_call_insn(insn):
         for ref in idautils.CodeRefsFrom(insn.ea, False):
@@ -320,28 +382,28 @@ def extract_function_calls_from(f, bb, insn):
 
 
 def extract_function_indirect_call_characteristic_features(f, bb, insn):
-    """ extract indirect function calls (e.g., call eax or call dword ptr [edx+4])
-        does not include calls like => call ds:dword_ABD4974
+    """extract indirect function calls (e.g., call eax or call dword ptr [edx+4])
+    does not include calls like => call ds:dword_ABD4974
 
-        most relevant at the function or basic block scope;
-        however, its most efficient to extract at the instruction scope
+    most relevant at the function or basic block scope;
+    however, its most efficient to extract at the instruction scope
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
     """
     if idaapi.is_call_insn(insn) and idc.get_operand_type(insn.ea, 0) in (idc.o_reg, idc.o_phrase, idc.o_displ):
         yield Characteristic("indirect call"), insn.ea
 
 
 def extract_features(f, bb, insn):
-    """ extract instruction features
+    """extract instruction features
 
-        args:
-            f (IDA func_t)
-            bb (IDA BasicBlock)
-            insn (IDA insn_t)
+    args:
+        f (IDA func_t)
+        bb (IDA BasicBlock)
+        insn (IDA insn_t)
     """
     for inst_handler in INSTRUCTION_HANDLERS:
         for (feature, ea) in inst_handler(f, bb, insn):

capa/features/extractors/loops.py

@@ -1,18 +1,24 @@
 # Copyright (C) 2020 FireEye, 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 networkx import nx
 from networkx.algorithms.components import strongly_connected_components
 
 
 def has_loop(edges, threshold=2):
-    """ check if a list of edges representing a directed graph contains a loop
+    """check if a list of edges representing a directed graph contains a loop
 
-        args:
-            edges: list of edge sets representing a directed graph i.e. [(1, 2), (2, 1)]
-            threshold: min number of nodes contained in loop
+    args:
+        edges: list of edge sets representing a directed graph i.e. [(1, 2), (2, 1)]
+        threshold: min number of nodes contained in loop
 
-        returns:
-            bool
+    returns:
+        bool
     """
     g = nx.DiGraph()
     g.add_edges_from(edges)

capa/features/extractors/strings.py

@@ -1,7 +1,12 @@
-# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
-#
 # strings code from FLOSS, https://github.com/fireeye/flare-floss
 #
+# Copyright (C) 2020 FireEye, 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 re
 from collections import namedtuple

capa/features/extractors/viv/__init__.py

@@ -1,13 +1,19 @@
 # Copyright (C) 2020 FireEye, 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 types
 
-import viv_utils
-
 import file
 import insn
 import function
+import viv_utils
 import basicblock
+
 import capa.features.extractors
 import capa.features.extractors.viv.file
 import capa.features.extractors.viv.insn

capa/features/extractors/viv/basicblock.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 string
 import struct

capa/features/extractors/viv/file.py

@@ -1,7 +1,14 @@
 # Copyright (C) 2020 FireEye, 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 PE.carve as pe_carve  # vivisect PE
 
+import capa.features.extractors.helpers
 import capa.features.extractors.strings
 from capa.features import String, Characteristic
 from capa.features.file import Export, Import, Section
@@ -35,11 +42,9 @@ def extract_file_import_names(vw, file_path):
         if is_viv_ord_impname(impname):
             # replace ord prefix with #
             impname = "#%s" % impname[len("ord") :]
-            tinfo = "%s.%s" % (modname, impname)
-            yield Import(tinfo), va
-        else:
-            yield Import(tinfo), va
-            yield Import(impname), va
+
+        for name in capa.features.extractors.helpers.generate_symbols(modname, impname):
+            yield Import(name), va
 
 
 def is_viv_ord_impname(impname):

capa/features/extractors/viv/function.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 vivisect.const
 
@@ -19,45 +25,6 @@ def interface_extract_function_XXX(f):
     yield NotImplementedError("feature"), NotImplementedError("virtual address")
 
 
-def get_switches(vw):
-    """
-    caching accessor to vivisect workspace switch constructs.
-    """
-    if "switches" in vw.metadata:
-        return vw.metadata["switches"]
-    else:
-        # addresses of switches in the program
-        switches = set()
-
-        for case_va, _ in filter(lambda t: "case" in t[1], vw.getNames()):
-            # assume that the xref to a case location is a switch construct
-            for switch_va, _, _, _ in vw.getXrefsTo(case_va):
-                switches.add(switch_va)
-
-        vw.metadata["switches"] = switches
-        return switches
-
-
-def get_functions_with_switch(vw):
-    if "functions_with_switch" in vw.metadata:
-        return vw.metadata["functions_with_switch"]
-    else:
-        functions = set()
-        for switch in get_switches(vw):
-            functions.add(vw.getFunction(switch))
-        vw.metadata["functions_with_switch"] = functions
-        return functions
-
-
-def extract_function_switch(f):
-    """
-    parse if a function contains a switch statement based on location names
-    method can be optimized
-    """
-    if f.va in get_functions_with_switch(f.vw):
-        yield Characteristic("switch"), f.va
-
-
 def extract_function_calls_to(f):
     for src, _, _, _ in f.vw.getXrefsTo(f.va, rtype=vivisect.const.REF_CODE):
         yield Characteristic("calls to"), src
@@ -72,7 +39,13 @@ def extract_function_loop(f):
     for bb in f.basic_blocks:
         if len(bb.instructions) > 0:
             for bva, bflags in bb.instructions[-1].getBranches():
-                if bflags & vivisect.envi.BR_COND or bflags & vivisect.envi.BR_FALL or bflags & vivisect.envi.BR_TABLE:
+                # vivisect does not set branch flags for non-conditional jmp so add explicit check
+                if (
+                    bflags & vivisect.envi.BR_COND
+                    or bflags & vivisect.envi.BR_FALL
+                    or bflags & vivisect.envi.BR_TABLE
+                    or bb.instructions[-1].mnem == "jmp"
+                ):
                     edges.append((bb.va, bva))
 
     if edges and loops.has_loop(edges):
@@ -94,4 +67,4 @@ def extract_features(f):
             yield feature, va
 
 
-FUNCTION_HANDLERS = (extract_function_switch, extract_function_calls_to, extract_function_loop)
+FUNCTION_HANDLERS = (extract_function_calls_to, extract_function_loop)

capa/features/extractors/viv/indirect_calls.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 collections
 

capa/features/extractors/viv/insn.py

@@ -1,12 +1,18 @@
 # Copyright (C) 2020 FireEye, 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 envi.memory
 import vivisect.const
 import envi.archs.i386.disasm
 
 import capa.features.extractors.helpers
-from capa.features import MAX_BYTES_FEATURE_SIZE, Bytes, String, Characteristic
-from capa.features.insn import Number, Offset, Mnemonic
+from capa.features import ARCH_X32, ARCH_X64, MAX_BYTES_FEATURE_SIZE, Bytes, String, Characteristic
+from capa.features.insn import API, Number, Offset, Mnemonic
 from capa.features.extractors.viv.indirect_calls import NotFoundError, resolve_indirect_call
 
 # security cookie checks may perform non-zeroing XORs, these are expected within a certain
@@ -14,6 +20,14 @@ from capa.features.extractors.viv.indirect_calls import NotFoundError, resolve_i
 SECURITY_COOKIE_BYTES_DELTA = 0x40
 
 
+def get_arch(vw):
+    arch = vw.getMeta("Architecture")
+    if arch == "i386":
+        return ARCH_X32
+    elif arch == "amd64":
+        return ARCH_X64
+
+
 def interface_extract_instruction_XXX(f, bb, insn):
     """
     parse features from the given instruction.
@@ -33,11 +47,15 @@ def get_imports(vw):
     """
     caching accessor to vivisect workspace imports
     avoids performance issues in vivisect when collecting locations
+
+    returns: Dict[int, Tuple[str, str]]
     """
     if "imports" in vw.metadata:
         return vw.metadata["imports"]
     else:
-        imports = {p[0]: p[3] for p in vw.getImports()}
+        imports = {
+            p[0]: (p[3].rpartition(".")[0], p[3].replace(".ord", ".#").rpartition(".")[2]) for p in vw.getImports()
+        }
         vw.metadata["imports"] = imports
         return imports
 
@@ -58,9 +76,10 @@ def extract_insn_api_features(f, bb, insn):
         target = oper.getOperAddr(insn)
 
         imports = get_imports(f.vw)
-        if target in imports.keys():
-            for feature, va in capa.features.extractors.helpers.generate_api_features(imports[target], insn.va):
-                yield feature, va
+        if target in imports:
+            dll, symbol = imports[target]
+            for name in capa.features.extractors.helpers.generate_symbols(dll, symbol):
+                yield API(name), insn.va
 
     # call via thunk on x86,
     # see 9324d1a8ae37a36ae560c37448c9705a at 0x407985
@@ -76,8 +95,11 @@ def extract_insn_api_features(f, bb, insn):
             return
         else:
             if thunk:
-                for feature, va in capa.features.extractors.helpers.generate_api_features(thunk, insn.va):
-                    yield feature, va
+                dll, _, symbol = thunk.rpartition(".")
+                if symbol.startswith("ord"):
+                    symbol = "#" + symbol[len("ord") :]
+                for name in capa.features.extractors.helpers.generate_symbols(dll, symbol):
+                    yield API(name), insn.va
 
     # call via import on x64
     # see Lab21-01.exe_:0x14000118C
@@ -86,9 +108,10 @@ def extract_insn_api_features(f, bb, insn):
         target = op.getOperAddr(insn)
 
         imports = get_imports(f.vw)
-        if target in imports.keys():
-            for feature, va in capa.features.extractors.helpers.generate_api_features(imports[target], insn.va):
-                yield feature, va
+        if target in imports:
+            dll, symbol = imports[target]
+            for name in capa.features.extractors.helpers.generate_symbols(dll, symbol):
+                yield API(name), insn.va
 
     elif isinstance(insn.opers[0], envi.archs.i386.disasm.i386RegOper):
         try:
@@ -102,9 +125,10 @@ def extract_insn_api_features(f, bb, insn):
             return
 
         imports = get_imports(f.vw)
-        if target in imports.keys():
-            for feature, va in capa.features.extractors.helpers.generate_api_features(imports[target], insn.va):
-                yield feature, va
+        if target in imports:
+            dll, symbol = imports[target]
+            for name in capa.features.extractors.helpers.generate_symbols(dll, symbol):
+                yield API(name), insn.va
 
 
 def extract_insn_number_features(f, bb, insn):
@@ -114,10 +138,13 @@ def extract_insn_number_features(f, bb, insn):
     #     push    3136B0h         ; dwControlCode
     for oper in insn.opers:
         # this is for both x32 and x64
-        if not isinstance(oper, envi.archs.i386.disasm.i386ImmOper):
+        if not isinstance(oper, (envi.archs.i386.disasm.i386ImmOper, envi.archs.i386.disasm.i386ImmMemOper)):
             continue
 
-        v = oper.getOperValue(oper)
+        if isinstance(oper, envi.archs.i386.disasm.i386ImmOper):
+            v = oper.getOperValue(oper)
+        else:
+            v = oper.getOperAddr(oper)
 
         if f.vw.probeMemory(v, 1, envi.memory.MM_READ):
             # this is a valid address
@@ -132,6 +159,77 @@ def extract_insn_number_features(f, bb, insn):
             return
 
         yield Number(v), insn.va
+        yield Number(v, arch=get_arch(f.vw)), insn.va
+
+
+def derefs(vw, p):
+    """
+    recursively follow the given pointer, yielding the valid memory addresses along the way.
+    useful when you may have a pointer to string, or pointer to pointer to string, etc.
+
+    this is a "do what i mean" type of helper function.
+    """
+    depth = 0
+    while True:
+        if not vw.isValidPointer(p):
+            return
+        yield p
+
+        try:
+            next = vw.readMemoryPtr(p)
+        except Exception:
+            # if not enough bytes can be read, such as end of the section.
+            # unfortunately, viv returns a plain old generic `Exception` for this.
+            return
+
+        # sanity: pointer points to self
+        if next == p:
+            return
+
+        # sanity: avoid chains of pointers that are unreasonably deep
+        depth += 1
+        if depth > 10:
+            return
+
+        p = next
+
+
+def read_memory(vw, va, size):
+    # as documented in #176, vivisect will not readMemory() when the section is not marked readable.
+    #
+    # but here, we don't care about permissions.
+    # so, copy the viv implementation of readMemory and remove the permissions check.
+    #
+    # this is derived from:
+    #   https://github.com/vivisect/vivisect/blob/5eb4d237bddd4069449a6bc094d332ceed6f9a96/envi/memory.py#L453-L462
+    for mva, mmaxva, mmap, mbytes in vw._map_defs:
+        if va >= mva and va < mmaxva:
+            mva, msize, mperms, mfname = mmap
+            offset = va - mva
+            return mbytes[offset : offset + size]
+    raise envi.SegmentationViolation(va)
+
+
+def read_bytes(vw, va):
+    """
+    read up to MAX_BYTES_FEATURE_SIZE from the given address.
+
+    raises:
+      envi.SegmentationViolation: if the given address is not valid.
+    """
+    segm = vw.getSegment(va)
+    if not segm:
+        raise envi.SegmentationViolation()
+
+    segm_end = segm[0] + segm[1]
+    try:
+        # Do not read beyond the end of a segment
+        if va + MAX_BYTES_FEATURE_SIZE > segm_end:
+            return read_memory(vw, va, segm_end - va)
+        else:
+            return read_memory(vw, va, MAX_BYTES_FEATURE_SIZE)
+    except envi.SegmentationViolation:
+        raise
 
 
 def extract_insn_bytes_features(f, bb, insn):
@@ -142,7 +240,6 @@ def extract_insn_bytes_features(f, bb, insn):
     """
     for oper in insn.opers:
         if insn.mnem == "call":
-            # ignore call instructions
             continue
 
         if isinstance(oper, envi.archs.i386.disasm.i386ImmOper):
@@ -151,28 +248,25 @@ def extract_insn_bytes_features(f, bb, insn):
             # handle case like:
             #   movzx   ecx, ds:byte_423258[eax]
             v = oper.disp
+        elif isinstance(oper, envi.archs.i386.disasm.i386SibOper):
+            # like 0x401000 in `mov eax, 0x401000[2 * ebx]`
+            v = oper.imm
         elif isinstance(oper, envi.archs.amd64.disasm.Amd64RipRelOper):
             # see: Lab21-01.exe_:0x1400010D3
             v = oper.getOperAddr(insn)
         else:
             continue
 
-        segm = f.vw.getSegment(v)
-        if not segm:
-            continue
+        for v in derefs(f.vw, v):
+            try:
+                buf = read_bytes(f.vw, v)
+            except envi.SegmentationViolation:
+                continue
 
-        segm_end = segm[0] + segm[1]
-        try:
-            # Do not read beyond the end of a segment
-            if v + MAX_BYTES_FEATURE_SIZE > segm_end:
-                extracted_bytes = f.vw.readMemory(v, segm_end - v)
-            else:
-                extracted_bytes = f.vw.readMemory(v, MAX_BYTES_FEATURE_SIZE)
-        except envi.SegmentationViolation:
-            pass
-        else:
-            if not capa.features.extractors.helpers.all_zeros(extracted_bytes):
-                yield Bytes(extracted_bytes), insn.va
+            if capa.features.extractors.helpers.all_zeros(buf):
+                continue
+
+            yield Bytes(buf), insn.va
 
 
 def read_string(vw, offset):
@@ -182,7 +276,7 @@ def read_string(vw, offset):
         pass
     else:
         if alen > 0:
-            return vw.readMemory(offset, alen).decode("utf-8")
+            return read_memory(vw, offset, alen).decode("utf-8")
 
     try:
         ulen = vw.detectUnicode(offset)
@@ -197,7 +291,7 @@ def read_string(vw, offset):
                 # vivisect seems to mis-detect the end unicode strings
                 # off by one, too short
                 ulen += 1
-            return vw.readMemory(offset, ulen).decode("utf-16")
+            return read_memory(vw, offset, ulen).decode("utf-16")
 
     raise ValueError("not a string", offset)
 
@@ -207,20 +301,25 @@ def extract_insn_string_features(f, bb, insn):
     # example:
     #
     #     push    offset aAcr     ; "ACR  > "
+
     for oper in insn.opers:
         if isinstance(oper, envi.archs.i386.disasm.i386ImmOper):
             v = oper.getOperValue(oper)
+        elif isinstance(oper, envi.archs.i386.disasm.i386SibOper):
+            # like 0x401000 in `mov eax, 0x401000[2 * ebx]`
+            v = oper.imm
         elif isinstance(oper, envi.archs.amd64.disasm.Amd64RipRelOper):
             v = oper.getOperAddr(insn)
         else:
             continue
 
-        try:
-            s = read_string(f.vw, v)
-        except ValueError:
-            continue
-        else:
-            yield String(s.rstrip("\x00")), insn.va
+        for v in derefs(f.vw, v):
+            try:
+                s = read_string(f.vw, v)
+            except ValueError:
+                continue
+            else:
+                yield String(s.rstrip("\x00")), insn.va
 
 
 def extract_insn_offset_features(f, bb, insn):
@@ -229,21 +328,38 @@ def extract_insn_offset_features(f, bb, insn):
     #
     #     .text:0040112F    cmp     [esi+4], ebx
     for oper in insn.opers:
+
         # this is for both x32 and x64
-        if not isinstance(oper, envi.archs.i386.disasm.i386RegMemOper):
-            continue
+        # like [esi + 4]
+        #       reg   ^
+        #             disp
+        if isinstance(oper, envi.archs.i386.disasm.i386RegMemOper):
+            if oper.reg == envi.archs.i386.disasm.REG_ESP:
+                continue
 
-        if oper.reg == envi.archs.i386.disasm.REG_ESP:
-            continue
+            if oper.reg == envi.archs.i386.disasm.REG_EBP:
+                continue
 
-        if oper.reg == envi.archs.i386.disasm.REG_EBP:
-            continue
+            # TODO: do x64 support for real.
+            if oper.reg == envi.archs.amd64.disasm.REG_RBP:
+                continue
 
-        # TODO: do x64 support for real.
-        if oper.reg == envi.archs.amd64.disasm.REG_RBP:
-            continue
+            # viv already decodes offsets as signed
+            v = oper.disp
 
-        yield Offset(oper.disp), insn.va
+            yield Offset(v), insn.va
+            yield Offset(v, arch=get_arch(f.vw)), insn.va
+
+        # like: [esi + ecx + 16384]
+        #        reg   ^     ^
+        #              index ^
+        #                    disp
+        elif isinstance(oper, envi.archs.i386.disasm.i386SibOper):
+            # viv already decodes offsets as signed
+            v = oper.disp
+
+            yield Offset(v), insn.va
+            yield Offset(v, arch=get_arch(f.vw)), insn.va
 
 
 def is_security_cookie(f, bb, insn):
@@ -305,7 +421,9 @@ def extract_insn_peb_access_characteristic_features(f, bb, insn):
     if insn.mnem not in ["push", "mov"]:
         return
 
-    if "fs" in insn.getPrefixName():
+    prefix = insn.getPrefixName()
+
+    if "fs" in prefix:
         for oper in insn.opers:
             # examples
             #
@@ -318,10 +436,12 @@ def extract_insn_peb_access_characteristic_features(f, bb, insn):
                 isinstance(oper, envi.archs.i386.disasm.i386ImmMemOper) and oper.imm == 0x30
             ):
                 yield Characteristic("peb access"), insn.va
-    elif "gs" in insn.getPrefixName():
+    elif "gs" in prefix:
         for oper in insn.opers:
-            if (isinstance(oper, envi.archs.amd64.disasm.i386RegMemOper) and oper.disp == 0x60) or (
-                isinstance(oper, envi.archs.amd64.disasm.i386ImmMemOper) and oper.imm == 0x60
+            if (
+                (isinstance(oper, envi.archs.amd64.disasm.i386RegMemOper) and oper.disp == 0x60)
+                or (isinstance(oper, envi.archs.amd64.disasm.i386SibOper) and oper.imm == 0x60)
+                or (isinstance(oper, envi.archs.amd64.disasm.i386ImmMemOper) and oper.imm == 0x60)
             ):
                 yield Characteristic("peb access"), insn.va
     else:

capa/features/file.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 capa.features import Feature
 
@@ -6,16 +12,16 @@ from capa.features import Feature
 class Export(Feature):
     def __init__(self, value, description=None):
         # value is export name
-        super(Export, self).__init__(value, description)
+        super(Export, self).__init__(value, description=description)
 
 
 class Import(Feature):
     def __init__(self, value, description=None):
         # value is import name
-        super(Import, self).__init__(value, description)
+        super(Import, self).__init__(value, description=description)
 
 
 class Section(Feature):
     def __init__(self, value, description=None):
         # value is section name
-        super(Section, self).__init__(value, description)
+        super(Section, self).__init__(value, description=description)

capa/features/freeze.py

@@ -41,6 +41,12 @@ json format:
     }
 
 Copyright (C) 2020 FireEye, 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 json
 import zlib
@@ -78,7 +84,16 @@ def dumps(extractor):
     returns:
       str: the serialized features.
     """
-    ret = {"version": 1, "functions": {}, "scopes": {"file": [], "function": [], "basic block": [], "instruction": [],}}
+    ret = {
+        "version": 1,
+        "functions": {},
+        "scopes": {
+            "file": [],
+            "function": [],
+            "basic block": [],
+            "instruction": [],
+        },
+    }
 
     for feature, va in extractor.extract_file_features():
         ret["scopes"]["file"].append(serialize_feature(feature) + (hex(va), ()))
@@ -93,14 +108,33 @@ def dumps(extractor):
             ret["functions"][hex(f)][hex(bb)] = []
 
             for feature, va in extractor.extract_basic_block_features(f, bb):
-                ret["scopes"]["basic block"].append(serialize_feature(feature) + (hex(va), (hex(f), hex(bb),)))
+                ret["scopes"]["basic block"].append(
+                    serialize_feature(feature)
+                    + (
+                        hex(va),
+                        (
+                            hex(f),
+                            hex(bb),
+                        ),
+                    )
+                )
 
-            for insn, insnva in sorted([(insn, int(insn)) for insn in extractor.get_instructions(f, bb)]):
+            for insnva, insn in sorted(
+                [(insn.__int__(), insn) for insn in extractor.get_instructions(f, bb)], key=lambda p: p[0]
+            ):
                 ret["functions"][hex(f)][hex(bb)].append(hex(insnva))
 
                 for feature, va in extractor.extract_insn_features(f, bb, insn):
                     ret["scopes"]["instruction"].append(
-                        serialize_feature(feature) + (hex(va), (hex(f), hex(bb), hex(insnva),))
+                        serialize_feature(feature)
+                        + (
+                            hex(va),
+                            (
+                                hex(f),
+                                hex(bb),
+                                hex(insnva),
+                            ),
+                        )
                     )
     return json.dumps(ret)
 
@@ -205,6 +239,7 @@ def load(buf):
 def main(argv=None):
     import sys
     import argparse
+
     import capa.main
 
     if argv is None:
@@ -238,12 +273,7 @@ def main(argv=None):
         logging.basicConfig(level=logging.INFO)
         logging.getLogger().setLevel(logging.INFO)
 
-    vw = capa.main.get_workspace(args.sample, args.format)
-
-    # don't import this at top level to support ida/py3 backend
-    import capa.features.extractors.viv
-
-    extractor = capa.features.extractors.viv.VivisectFeatureExtractor(vw, args.sample)
+    extractor = capa.main.get_extractor(args.sample, args.format)
     with open(args.output, "wb") as f:
         f.write(dump(extractor))
 

capa/features/insn.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 capa.features import Feature
 
@@ -14,16 +20,16 @@ class API(Feature):
 
 
 class Number(Feature):
-    def __init__(self, value, description=None):
-        super(Number, self).__init__(value, description)
+    def __init__(self, value, arch=None, description=None):
+        super(Number, self).__init__(value, arch=arch, description=description)
 
     def get_value_str(self):
         return "0x%X" % self.value
 
 
 class Offset(Feature):
-    def __init__(self, value, description=None):
-        super(Offset, self).__init__(value, description)
+    def __init__(self, value, arch=None, description=None):
+        super(Offset, self).__init__(value, arch=arch, description=description)
 
     def get_value_str(self):
         return "0x%X" % self.value
@@ -31,4 +37,4 @@ class Offset(Feature):
 
 class Mnemonic(Feature):
     def __init__(self, value, description=None):
-        super(Mnemonic, self).__init__(value, description)
+        super(Mnemonic, self).__init__(value, description=description)

capa/helpers.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 os
 

capa/ida/explorer/proxy.py

@@ -1,83 +0,0 @@
-# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
-
-from PyQt5 import QtCore
-
-from capa.ida.explorer.model import CapaExplorerDataModel
-
-
-class CapaExplorerSortFilterProxyModel(QtCore.QSortFilterProxyModel):
-    def __init__(self, parent=None):
-        """ """
-        super(CapaExplorerSortFilterProxyModel, self).__init__(parent)
-
-    def lessThan(self, left, right):
-        """ true if the value of the left item is less than value of right item
-
-            @param left: QModelIndex*
-            @param right: QModelIndex*
-
-            @retval True/False
-        """
-        ldata = left.internalPointer().data(left.column())
-        rdata = right.internalPointer().data(right.column())
-
-        if (
-            ldata
-            and rdata
-            and left.column() == CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS
-            and left.column() == right.column()
-        ):
-            # convert virtual address before compare
-            return int(ldata, 16) < int(rdata, 16)
-        else:
-            # compare as lowercase
-            return ldata.lower() < rdata.lower()
-
-    def filterAcceptsRow(self, row, parent):
-        """ true if the item in the row indicated by the given row and parent
-            should be included in the model; otherwise returns false
-
-            @param row: int
-            @param parent: QModelIndex*
-
-            @retval True/False
-        """
-        if self.filter_accepts_row_self(row, parent):
-            return True
-
-        alpha = parent
-        while alpha.isValid():
-            if self.filter_accepts_row_self(alpha.row(), alpha.parent()):
-                return True
-            alpha = alpha.parent()
-
-        if self.index_has_accepted_children(row, parent):
-            return True
-
-        return False
-
-    def add_single_string_filter(self, column, string):
-        """ add fixed string filter
-
-            @param column: key column
-            @param string: string to sort
-        """
-        self.setFilterKeyColumn(column)
-        self.setFilterFixedString(string)
-
-    def index_has_accepted_children(self, row, parent):
-        """ """
-        model_index = self.sourceModel().index(row, 0, parent)
-
-        if model_index.isValid():
-            for idx in range(self.sourceModel().rowCount(model_index)):
-                if self.filter_accepts_row_self(idx, model_index):
-                    return True
-                if self.index_has_accepted_children(idx, model_index):
-                    return True
-
-        return False
-
-    def filter_accepts_row_self(self, row, parent):
-        """ """
-        return super(CapaExplorerSortFilterProxyModel, self).filterAcceptsRow(row, parent)

capa/ida/explorer/view.py

@@ -1,256 +0,0 @@
-# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
-
-import idc
-import idaapi
-from PyQt5 import QtGui, QtCore, QtWidgets
-
-from capa.ida.explorer.item import CapaExplorerRuleItem, CapaExplorerFunctionItem
-from capa.ida.explorer.model import CapaExplorerDataModel
-
-
-class CapaExplorerQtreeView(QtWidgets.QTreeView):
-    """ capa explorer QTreeView implementation
-
-        view controls UI action responses and displays data from
-        CapaExplorerDataModel
-
-        view does not modify CapaExplorerDataModel directly - data
-        modifications should be implemented in CapaExplorerDataModel
-    """
-
-    def __init__(self, model, parent=None):
-        """ initialize CapaExplorerQTreeView """
-        super(CapaExplorerQtreeView, self).__init__(parent)
-
-        self.setModel(model)
-
-        self.model = model
-        self.parent = parent
-
-        # configure custom UI controls
-        self.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
-        self.setExpandsOnDoubleClick(False)
-        self.setSortingEnabled(True)
-        self.model.setDynamicSortFilter(False)
-
-        # configure view columns to auto-resize
-        for idx in range(CapaExplorerDataModel.COLUMN_COUNT):
-            self.header().setSectionResizeMode(idx, QtWidgets.QHeaderView.Interactive)
-
-        # connect slots to resize columns when expanded or collapsed
-        self.expanded.connect(self.resize_columns_to_content)
-        self.collapsed.connect(self.resize_columns_to_content)
-
-        # connect slots
-        self.customContextMenuRequested.connect(self.slot_custom_context_menu_requested)
-        self.doubleClicked.connect(self.slot_double_click)
-
-        self.setStyleSheet("QTreeView::item {padding-right: 15 px;padding-bottom: 2 px;}")
-
-    def reset(self):
-        """ reset user interface changes
-
-            called when view should reset any user interface changes
-            made since the last reset e.g. IDA window highlighting
-        """
-        self.collapseAll()
-        self.resize_columns_to_content()
-
-    def resize_columns_to_content(self):
-        """ reset view columns to contents """
-        self.header().resizeSections(QtWidgets.QHeaderView.ResizeToContents)
-
-    def map_index_to_source_item(self, model_index):
-        """ map proxy model index to source model item
-
-            @param model_index: QModelIndex*
-
-            @retval QObject*
-        """
-        return self.model.mapToSource(model_index).internalPointer()
-
-    def send_data_to_clipboard(self, data):
-        """ copy data to the clipboard
-
-            @param data: data to be copied
-        """
-        clip = QtWidgets.QApplication.clipboard()
-        clip.clear(mode=clip.Clipboard)
-        clip.setText(data, mode=clip.Clipboard)
-
-    def new_action(self, display, data, slot):
-        """ create action for context menu
-
-            @param display: text displayed to user in context menu
-            @param data: data passed to slot
-            @param slot: slot to connect
-
-            @retval QAction*
-        """
-        action = QtWidgets.QAction(display, self.parent)
-        action.setData(data)
-        action.triggered.connect(lambda checked: slot(action))
-
-        return action
-
-    def load_default_context_menu_actions(self, data):
-        """ yield actions specific to function custom context menu
-
-            @param data: tuple
-
-            @yield QAction*
-        """
-        default_actions = (
-            ("Copy column", data, self.slot_copy_column),
-            ("Copy row", data, self.slot_copy_row),
-        )
-
-        # add default actions
-        for action in default_actions:
-            yield self.new_action(*action)
-
-    def load_function_context_menu_actions(self, data):
-        """ yield actions specific to function custom context menu
-
-            @param data: tuple
-
-            @yield QAction*
-        """
-        function_actions = (("Rename function", data, self.slot_rename_function),)
-
-        # add function actions
-        for action in function_actions:
-            yield self.new_action(*action)
-
-        # add default actions
-        for action in self.load_default_context_menu_actions(data):
-            yield action
-
-    def load_default_context_menu(self, pos, item, model_index):
-        """ create default custom context menu
-
-            creates custom context menu containing default actions
-
-            @param pos: TODO
-            @param item: TODO
-            @param model_index: TODO
-
-            @retval QMenu*
-        """
-        menu = QtWidgets.QMenu()
-
-        for action in self.load_default_context_menu_actions((pos, item, model_index)):
-            menu.addAction(action)
-
-        return menu
-
-    def load_function_item_context_menu(self, pos, item, model_index):
-        """ create function custom context menu
-
-            creates custom context menu containing actions specific to functions
-            and the default actions
-
-            @param pos: TODO
-            @param item: TODO
-            @param model_index: TODO
-
-            @retval QMenu*
-        """
-        menu = QtWidgets.QMenu()
-
-        for action in self.load_function_context_menu_actions((pos, item, model_index)):
-            menu.addAction(action)
-
-        return menu
-
-    def show_custom_context_menu(self, menu, pos):
-        """ display custom context menu in view
-
-            @param menu: TODO
-            @param pos: TODO
-        """
-        if menu:
-            menu.exec_(self.viewport().mapToGlobal(pos))
-
-    def slot_copy_column(self, action):
-        """ slot connected to custom context menu
-
-            allows user to select a column and copy the data
-            to clipboard
-
-            @param action: QAction*
-        """
-        _, item, model_index = action.data()
-        self.send_data_to_clipboard(item.data(model_index.column()))
-
-    def slot_copy_row(self, action):
-        """ slot connected to custom context menu
-
-            allows user to select a row and copy the space-delimited
-            data to clipboard
-
-            @param action: QAction*
-        """
-        _, item, _ = action.data()
-        self.send_data_to_clipboard(str(item))
-
-    def slot_rename_function(self, action):
-        """ slot connected to custom context menu
-
-            allows user to select a edit a function name and push
-            changes to IDA
-
-            @param action: QAction*
-        """
-        _, item, model_index = action.data()
-
-        # make item temporary edit, reset after user is finished
-        item.setIsEditable(True)
-        self.edit(model_index)
-        item.setIsEditable(False)
-
-    def slot_custom_context_menu_requested(self, pos):
-        """ slot connected to custom context menu request
-
-            displays custom context menu to user containing action
-            relevant to the data item selected
-
-            @param pos: TODO
-        """
-        model_index = self.indexAt(pos)
-
-        if not model_index.isValid():
-            return
-
-        item = self.map_index_to_source_item(model_index)
-        column = model_index.column()
-        menu = None
-
-        if CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION == column and isinstance(item, CapaExplorerFunctionItem):
-            # user hovered function item
-            menu = self.load_function_item_context_menu(pos, item, model_index)
-        else:
-            # user hovered default item
-            menu = self.load_default_context_menu(pos, item, model_index)
-
-        # show custom context menu at view position
-        self.show_custom_context_menu(menu, pos)
-
-    def slot_double_click(self, model_index):
-        """ slot connected to double click event
-
-            @param model_index: QModelIndex*
-        """
-        if not model_index.isValid():
-            return
-
-        item = self.map_index_to_source_item(model_index)
-        column = model_index.column()
-
-        if CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS == column and item.location:
-            # user double-clicked virtual address column - navigate IDA to address
-            idc.jumpto(item.location)
-
-        if CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION == column:
-            # user double-clicked information column - un/expand
-            self.collapse(model_index) if self.isExpanded(model_index) else self.expand(model_index)

capa/ida/helpers/__init__.py

@@ -1,9 +1,16 @@
 # Copyright (C) 2020 FireEye, 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
 import datetime
 
 import idc
+import six
 import idaapi
 import idautils
 
@@ -11,6 +18,14 @@ import capa
 
 logger = logging.getLogger("capa")
 
+SUPPORTED_IDA_VERSIONS = [
+    "7.1",
+    "7.2",
+    "7.3",
+    "7.4",
+    "7.5",
+]
+
 # file type names as returned by idaapi.get_file_type_name()
 SUPPORTED_FILE_TYPES = [
     "Portable executable for 80386 (PE)",
@@ -23,6 +38,18 @@ def inform_user_ida_ui(message):
     idaapi.info("%s. Please refer to IDA Output window for more information." % message)
 
 
+def is_supported_ida_version():
+    version = idaapi.get_kernel_version()
+    if version not in SUPPORTED_IDA_VERSIONS:
+        warning_msg = "This plugin does not support your IDA Pro version"
+        logger.warning(warning_msg)
+        logger.warning(
+            "Your IDA Pro version is: %s. Supported versions are: %s." % (version, ", ".join(SUPPORTED_IDA_VERSIONS))
+        )
+        return False
+    return True
+
+
 def is_supported_file_type():
     file_type = idaapi.get_file_type_name()
     if file_type not in SUPPORTED_FILE_TYPES:
@@ -34,7 +61,6 @@ def is_supported_file_type():
         )
         logger.error(" If you don't know the input file type, you can try using the `file` utility to guess it.")
         logger.error("-" * 80)
-        inform_user_ida_ui("capa does not support the format of this file")
         return False
     return True
 
@@ -57,15 +83,26 @@ def get_func_start_ea(ea):
 
 
 def collect_metadata():
+    md5 = idautils.GetInputFileMD5()
+    if not isinstance(md5, six.string_types):
+        md5 = capa.features.bytes_to_str(md5)
+
+    sha256 = idaapi.retrieve_input_file_sha256()
+    if not isinstance(sha256, six.string_types):
+        sha256 = capa.features.bytes_to_str(sha256)
+
     return {
         "timestamp": datetime.datetime.now().isoformat(),
         # "argv" is not relevant here
         "sample": {
-            "md5": capa.features.bytes_to_str(idautils.GetInputFileMD5()),
-            # "sha1" not easily accessible
-            "sha256": capa.features.bytes_to_str(idaapi.retrieve_input_file_sha256()),
+            "md5": md5,
+            "sha1": "",  # not easily accessible
+            "sha256": sha256,
             "path": idaapi.get_input_file_path(),
         },
-        "analysis": {"format": idaapi.get_file_type_name(), "extractor": "ida",},
+        "analysis": {
+            "format": idaapi.get_file_type_name(),
+            "extractor": "ida",
+        },
         "version": capa.version.__version__,
     }

capa/ida/ida_capa_explorer.py

@@ -1,564 +0,0 @@
-# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
-
-import os
-import json
-import logging
-import collections
-
-import idaapi
-from PyQt5 import QtGui, QtCore, QtWidgets
-
-import capa.main
-import capa.rules
-import capa.ida.helpers
-import capa.render.utils as rutils
-import capa.features.extractors.ida
-from capa.ida.explorer.view import CapaExplorerQtreeView
-from capa.ida.explorer.model import CapaExplorerDataModel
-from capa.ida.explorer.proxy import CapaExplorerSortFilterProxyModel
-
-PLUGIN_NAME = "capa explorer"
-
-logger = logging.getLogger("capa")
-
-
-class CapaExplorerIdaHooks(idaapi.UI_Hooks):
-    def __init__(self, screen_ea_changed_hook, action_hooks):
-        """ facilitate IDA UI hooks
-
-            @param screen_ea_changed_hook: function hook for IDA screen ea changed
-            @param action_hooks: dict of IDA action handles
-        """
-        super(CapaExplorerIdaHooks, self).__init__()
-
-        self.screen_ea_changed_hook = screen_ea_changed_hook
-        self.process_action_hooks = action_hooks
-        self.process_action_handle = None
-        self.process_action_meta = {}
-
-    def preprocess_action(self, name):
-        """ called prior to action completed
-
-            @param name: name of action defined by idagui.cfg
-
-            @retval must be 0
-        """
-        self.process_action_handle = self.process_action_hooks.get(name, None)
-
-        if self.process_action_handle:
-            self.process_action_handle(self.process_action_meta)
-
-        # must return 0 for IDA
-        return 0
-
-    def postprocess_action(self):
-        """ called after action completed """
-        if not self.process_action_handle:
-            return
-
-        self.process_action_handle(self.process_action_meta, post=True)
-        self.reset()
-
-    def screen_ea_changed(self, curr_ea, prev_ea):
-        """ called after screen location is changed
-
-            @param curr_ea: current location
-            @param prev_ea: prev location
-        """
-        self.screen_ea_changed_hook(idaapi.get_current_widget(), curr_ea, prev_ea)
-
-    def reset(self):
-        """ reset internal state """
-        self.process_action_handle = None
-        self.process_action_meta.clear()
-
-
-class CapaExplorerForm(idaapi.PluginForm):
-    def __init__(self):
-        """ """
-        super(CapaExplorerForm, self).__init__()
-
-        self.form_title = PLUGIN_NAME
-        self.file_loc = __file__
-
-        self.parent = None
-        self.ida_hooks = None
-        self.doc = None
-
-        # models
-        self.model_data = None
-        self.model_proxy = None
-
-        # user interface elements
-        self.view_limit_results_by_function = None
-        self.view_tree = None
-        self.view_summary = None
-        self.view_attack = None
-        self.view_tabs = None
-        self.view_menu_bar = None
-
-    def OnCreate(self, form):
-        """ """
-        self.parent = self.FormToPyQtWidget(form)
-        self.load_interface()
-        self.load_capa_results()
-        self.load_ida_hooks()
-
-        self.view_tree.reset()
-
-        logger.info("form created.")
-
-    def Show(self):
-        """ """
-        return idaapi.PluginForm.Show(
-            self, self.form_title, options=(idaapi.PluginForm.WOPN_TAB | idaapi.PluginForm.WCLS_CLOSE_LATER)
-        )
-
-    def OnClose(self, form):
-        """ form is closed """
-        self.unload_ida_hooks()
-        self.ida_reset()
-
-        logger.info("form closed.")
-
-    def load_interface(self):
-        """ load user interface """
-        # load models
-        self.model_data = CapaExplorerDataModel()
-        self.model_proxy = CapaExplorerSortFilterProxyModel()
-        self.model_proxy.setSourceModel(self.model_data)
-
-        # load tree
-        self.view_tree = CapaExplorerQtreeView(self.model_proxy, self.parent)
-
-        # load summary table
-        self.load_view_summary()
-        self.load_view_attack()
-
-        # load parent tab and children tab views
-        self.load_view_tabs()
-        self.load_view_checkbox_limit_by()
-        self.load_view_summary_tab()
-        self.load_view_attack_tab()
-        self.load_view_tree_tab()
-
-        # load menu bar and sub menus
-        self.load_view_menu_bar()
-        self.load_file_menu()
-
-        # load parent view
-        self.load_view_parent()
-
-    def load_view_tabs(self):
-        """ load tabs """
-        tabs = QtWidgets.QTabWidget()
-        self.view_tabs = tabs
-
-    def load_view_menu_bar(self):
-        """ load menu bar """
-        bar = QtWidgets.QMenuBar()
-        self.view_menu_bar = bar
-
-    def load_view_summary(self):
-        """ load capa summary table """
-        table_headers = [
-            "Capability",
-            "Namespace",
-        ]
-
-        table = QtWidgets.QTableWidget()
-
-        table.setColumnCount(len(table_headers))
-        table.verticalHeader().setVisible(False)
-        table.setSortingEnabled(False)
-        table.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
-        table.setFocusPolicy(QtCore.Qt.NoFocus)
-        table.setSelectionMode(QtWidgets.QAbstractItemView.NoSelection)
-        table.setHorizontalHeaderLabels(table_headers)
-        table.horizontalHeader().setDefaultAlignment(QtCore.Qt.AlignLeft)
-        table.setShowGrid(False)
-        table.setStyleSheet("QTableWidget::item { padding: 25px; }")
-
-        self.view_summary = table
-
-    def load_view_attack(self):
-        """ load MITRE ATT&CK table """
-        table_headers = [
-            "ATT&CK Tactic",
-            "ATT&CK Technique ",
-        ]
-
-        table = QtWidgets.QTableWidget()
-
-        table.setColumnCount(len(table_headers))
-        table.verticalHeader().setVisible(False)
-        table.setSortingEnabled(False)
-        table.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
-        table.setFocusPolicy(QtCore.Qt.NoFocus)
-        table.setSelectionMode(QtWidgets.QAbstractItemView.NoSelection)
-        table.setHorizontalHeaderLabels(table_headers)
-        table.horizontalHeader().setDefaultAlignment(QtCore.Qt.AlignLeft)
-        table.setShowGrid(False)
-        table.setStyleSheet("QTableWidget::item { padding: 25px; }")
-
-        self.view_attack = table
-
-    def load_view_checkbox_limit_by(self):
-        """ load limit results by function checkbox """
-        check = QtWidgets.QCheckBox("Limit results to current function")
-        check.setChecked(False)
-        check.stateChanged.connect(self.slot_checkbox_limit_by_changed)
-
-        self.view_limit_results_by_function = check
-
-    def load_view_parent(self):
-        """ load view parent """
-        layout = QtWidgets.QVBoxLayout()
-
-        layout.addWidget(self.view_tabs)
-        layout.setMenuBar(self.view_menu_bar)
-
-        self.parent.setLayout(layout)
-
-    def load_view_tree_tab(self):
-        """ load capa tree tab view """
-        layout = QtWidgets.QVBoxLayout()
-        layout.addWidget(self.view_limit_results_by_function)
-        layout.addWidget(self.view_tree)
-
-        tab = QtWidgets.QWidget()
-        tab.setLayout(layout)
-
-        self.view_tabs.addTab(tab, "Tree View")
-
-    def load_view_summary_tab(self):
-        """ load capa summary tab view """
-        layout = QtWidgets.QVBoxLayout()
-        layout.addWidget(self.view_summary)
-
-        tab = QtWidgets.QWidget()
-        tab.setLayout(layout)
-
-        self.view_tabs.addTab(tab, "Summary")
-
-    def load_view_attack_tab(self):
-        """ load MITRE ATT&CK tab view """
-        layout = QtWidgets.QVBoxLayout()
-        layout.addWidget(self.view_attack)
-
-        tab = QtWidgets.QWidget()
-        tab.setLayout(layout)
-
-        self.view_tabs.addTab(tab, "MITRE")
-
-    def load_file_menu(self):
-        """ load file menu actions """
-        actions = (
-            ("Reset view", "Reset plugin view", self.reset),
-            ("Run analysis", "Run capa analysis on current database", self.reload),
-            ("Export results...", "Export capa results as JSON file", self.export_json),
-        )
-
-        menu = self.view_menu_bar.addMenu("File")
-        for (name, _, handle) in actions:
-            action = QtWidgets.QAction(name, self.parent)
-            action.triggered.connect(handle)
-            menu.addAction(action)
-
-    def export_json(self):
-        """ export capa results as JSON file """
-        if not self.doc:
-            idaapi.info("No capa results to export.")
-            return
-        path = idaapi.ask_file(True, "*.json", "Choose file")
-        if os.path.exists(path) and 1 != idaapi.ask_yn(1, "File already exists. Overwrite?"):
-            return
-        with open(path, "wb") as export_file:
-            export_file.write(
-                json.dumps(self.doc, sort_keys=True, cls=capa.render.CapaJsonObjectEncoder).encode("utf-8")
-            )
-
-    def load_ida_hooks(self):
-        """ load IDA Pro UI hooks """
-        action_hooks = {
-            "MakeName": self.ida_hook_rename,
-            "EditFunction": self.ida_hook_rename,
-        }
-
-        self.ida_hooks = CapaExplorerIdaHooks(self.ida_hook_screen_ea_changed, action_hooks)
-        self.ida_hooks.hook()
-
-    def unload_ida_hooks(self):
-        """ unload IDA Pro UI hooks """
-        if self.ida_hooks:
-            self.ida_hooks.unhook()
-
-    def ida_hook_rename(self, meta, post=False):
-        """ hook for IDA rename action
-
-            called twice, once before action and once after
-            action completes
-
-            @param meta: metadata cache
-            @param post: indicates pre or post action
-        """
-        location = idaapi.get_screen_ea()
-        if not location or not capa.ida.helpers.is_func_start(location):
-            return
-
-        curr_name = idaapi.get_name(location)
-
-        if post:
-            # post action update data model w/ current name
-            self.model_data.update_function_name(meta.get("prev_name", ""), curr_name)
-        else:
-            # pre action so save current name for replacement later
-            meta["prev_name"] = curr_name
-
-    def ida_hook_screen_ea_changed(self, widget, new_ea, old_ea):
-        """ hook for IDA screen ea changed
-
-            @param widget: IDA widget type
-            @param new_ea: destination ea
-            @param old_ea: source ea
-         """
-        if not self.view_limit_results_by_function.isChecked():
-            # ignore if checkbox not selected
-            return
-
-        if idaapi.get_widget_type(widget) != idaapi.BWN_DISASM:
-            # ignore views other than asm
-            return
-
-        # attempt to map virtual addresses to function start addresses
-        new_func_start = capa.ida.helpers.get_func_start_ea(new_ea)
-        old_func_start = capa.ida.helpers.get_func_start_ea(old_ea)
-
-        if new_func_start and new_func_start == old_func_start:
-            # navigated within the same function - do nothing
-            return
-
-        if new_func_start:
-            # navigated to new function - filter for function start virtual address
-            match = capa.ida.explorer.item.location_to_hex(new_func_start)
-        else:
-            # navigated to virtual address not in valid function - clear filter
-            match = ""
-
-        # filter on virtual address to avoid updating filter string if function name is changed
-        self.model_proxy.add_single_string_filter(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS, match)
-        self.view_tree.resize_columns_to_content()
-
-    def load_capa_results(self):
-        """ run capa analysis and render results in UI """
-        logger.info("-" * 80)
-        logger.info(" Using default embedded rules.")
-        logger.info(" ")
-        logger.info(" You can see the current default rule set here:")
-        logger.info("     https://github.com/fireeye/capa-rules")
-        logger.info("-" * 80)
-
-        rules_path = os.path.join(os.path.dirname(self.file_loc), "../..", "rules")
-        rules = capa.main.get_rules(rules_path)
-        rules = capa.rules.RuleSet(rules)
-
-        meta = capa.ida.helpers.collect_metadata()
-
-        capabilities, counts = capa.main.find_capabilities(
-            rules, capa.features.extractors.ida.IdaFeatureExtractor(), True
-        )
-        meta["analysis"].update(counts)
-
-        # support binary files specifically for x86/AMD64 shellcode
-        # warn user binary file is loaded but still allow capa to process it
-        # TODO: check specific architecture of binary files based on how user configured IDA processors
-        if idaapi.get_file_type_name() == "Binary file":
-            logger.warning("-" * 80)
-            logger.warning(" Input file appears to be a binary file.")
-            logger.warning(" ")
-            logger.warning(
-                " capa currently only supports analyzing binary files containing x86/AMD64 shellcode with IDA."
-            )
-            logger.warning(
-                " This means the results may be misleading or incomplete if the binary file loaded in IDA is not x86/AMD64."
-            )
-            logger.warning(" If you don't know the input file type, you can try using the `file` utility to guess it.")
-            logger.warning("-" * 80)
-
-            capa.ida.helpers.inform_user_ida_ui("capa encountered warnings during analysis")
-
-        if capa.main.has_file_limitation(rules, capabilities, is_standalone=False):
-            capa.ida.helpers.inform_user_ida_ui("capa encountered warnings during analysis")
-
-        logger.info("analysis completed.")
-
-        self.doc = capa.render.convert_capabilities_to_result_document(meta, rules, capabilities)
-
-        self.model_data.render_capa_doc(self.doc)
-        self.render_capa_doc_summary()
-        self.render_capa_doc_mitre_summary()
-
-        self.set_view_tree_default_sort_order()
-
-        logger.info("render views completed.")
-
-    def set_view_tree_default_sort_order(self):
-        """ set capa tree view default sort order """
-        self.view_tree.sortByColumn(CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION, QtCore.Qt.AscendingOrder)
-
-    def render_capa_doc_summary(self):
-        """ render capa summary results """
-        for (row, rule) in enumerate(rutils.capability_rules(self.doc)):
-            count = len(rule["matches"])
-
-            if count == 1:
-                capability = rule["meta"]["name"]
-            else:
-                capability = "%s (%d matches)" % (rule["meta"]["name"], count)
-
-            self.view_summary.setRowCount(row + 1)
-
-            self.view_summary.setItem(row, 0, self.render_new_table_header_item(capability))
-            self.view_summary.setItem(row, 1, QtWidgets.QTableWidgetItem(rule["meta"]["namespace"]))
-
-        # resize columns to content
-        self.view_summary.resizeColumnsToContents()
-
-    def render_capa_doc_mitre_summary(self):
-        """ render capa MITRE ATT&CK results """
-        tactics = collections.defaultdict(set)
-
-        for rule in rutils.capability_rules(self.doc):
-            if not rule["meta"].get("att&ck"):
-                continue
-
-            for attack in rule["meta"]["att&ck"]:
-                tactic, _, rest = attack.partition("::")
-                if "::" in rest:
-                    technique, _, rest = rest.partition("::")
-                    subtechnique, _, id = rest.rpartition(" ")
-                    tactics[tactic].add((technique, subtechnique, id))
-                else:
-                    technique, _, id = rest.rpartition(" ")
-                    tactics[tactic].add((technique, id))
-
-        column_one = []
-        column_two = []
-
-        for (tactic, techniques) in sorted(tactics.items()):
-            column_one.append(tactic.upper())
-            # add extra space when more than one technique
-            column_one.extend(["" for i in range(len(techniques) - 1)])
-
-            for spec in sorted(techniques):
-                if len(spec) == 2:
-                    technique, id = spec
-                    column_two.append("%s %s" % (technique, id))
-                elif len(spec) == 3:
-                    technique, subtechnique, id = spec
-                    column_two.append("%s::%s %s" % (technique, subtechnique, id))
-                else:
-                    raise RuntimeError("unexpected ATT&CK spec format")
-
-        self.view_attack.setRowCount(max(len(column_one), len(column_two)))
-
-        for row, value in enumerate(column_one):
-            self.view_attack.setItem(row, 0, self.render_new_table_header_item(value))
-
-        for row, value in enumerate(column_two):
-            self.view_attack.setItem(row, 1, QtWidgets.QTableWidgetItem(value))
-
-        # resize columns to content
-        self.view_attack.resizeColumnsToContents()
-
-    def render_new_table_header_item(self, text):
-        """ create new table header item with default style """
-        item = QtWidgets.QTableWidgetItem(text)
-        item.setForeground(QtGui.QColor(88, 139, 174))
-
-        font = QtGui.QFont()
-        font.setBold(True)
-
-        item.setFont(font)
-
-        return item
-
-    def ida_reset(self):
-        """ reset IDA UI """
-        self.model_data.reset()
-        self.view_tree.reset()
-        self.view_limit_results_by_function.setChecked(False)
-        self.set_view_tree_default_sort_order()
-
-    def reload(self):
-        """ reload views and re-run capa analysis """
-        self.ida_reset()
-        self.model_proxy.invalidate()
-        self.model_data.clear()
-        self.view_summary.setRowCount(0)
-        self.load_capa_results()
-
-        logger.info("reload complete.")
-        idaapi.info("%s reload completed." % PLUGIN_NAME)
-
-    def reset(self):
-        """ reset UI elements
-
-            e.g. checkboxes and IDA highlighting
-        """
-        self.ida_reset()
-
-        logger.info("reset completed.")
-        idaapi.info("%s reset completed." % PLUGIN_NAME)
-
-    def slot_menu_bar_hovered(self, action):
-        """ display menu action tooltip
-
-            @param action: QtWidgets.QAction*
-
-            @reference: https://stackoverflow.com/questions/21725119/why-wont-qtooltips-appear-on-qactions-within-a-qmenu
-        """
-        QtWidgets.QToolTip.showText(
-            QtGui.QCursor.pos(), action.toolTip(), self.view_menu_bar, self.view_menu_bar.actionGeometry(action)
-        )
-
-    def slot_checkbox_limit_by_changed(self):
-        """ slot activated if checkbox clicked
-
-            if checked, configure function filter if screen location is located
-            in function, otherwise clear filter
-        """
-        match = ""
-        if self.view_limit_results_by_function.isChecked():
-            location = capa.ida.helpers.get_func_start_ea(idaapi.get_screen_ea())
-            if location:
-                match = capa.ida.explorer.item.location_to_hex(location)
-
-        self.model_proxy.add_single_string_filter(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS, match)
-
-        self.view_tree.resize_columns_to_content()
-
-
-def main():
-    """ TODO: move to idaapi.plugin_t class """
-    logging.basicConfig(level=logging.INFO)
-
-    if not capa.ida.helpers.is_supported_file_type():
-        return -1
-
-    global CAPA_EXPLORER_FORM
-
-    try:
-        # there is an instance, reload it
-        CAPA_EXPLORER_FORM
-        CAPA_EXPLORER_FORM.Close()
-        CAPA_EXPLORER_FORM = CapaExplorerForm()
-    except Exception:
-        # there is no instance yet
-        CAPA_EXPLORER_FORM = CapaExplorerForm()
-
-    CAPA_EXPLORER_FORM.Show()
-
-
-if __name__ == "__main__":
-    main()

capa/ida/plugin/README.md

@@ -0,0 +1,111 @@
+# capa explorer
+
+capa explorer is an IDA Pro plugin that integrates the FLARE team's open-source framework, capa, with IDA. capa is a framework that uses a well-defined collection of rules to 
+identify capabilities in a program. You can run capa against a PE file or shellcode and it tells you what it thinks the program can do. For example, it might suggest that 
+the program is a backdoor, can install services, or relies on HTTP to communicate. 
+
+The capa explorer IDA plugin brings capa's detection capabilities to IDA. You can use capa explorer to run capa directly on an IDA database without needing access
+to the source binary. Once a database has been analyzed, capa explorer can be used to quickly identify and navigate to interesting areas of a program 
+and dissect capa rule matches at the assembly level.
+
+To illustrate, we use capa explorer to analyze Lab 14-02 from [Practical Malware Analysis](https://nostarch.com/malware) (PMA) available [here](https://practicalmalwareanalysis.com/labs/). Our
+goal is to understand the program's functionality.
+
+After loading Lab 14-02 into IDA and analyzing the database with capa explorer, we see that capa detected a rule match for `self delete via COMSPEC environment variable`:
+
+![](../../../doc/img/ida_plugin_example_1.png)
+
+We can use capa explorer to navigate the IDA Disassembly view directly to the suspect function and get an assembly-level breakdown of why capa matched `self delete via COMSPEC environment variable` 
+for this particular function.
+
+![](../../../doc/img/ida_plugin_example_2.png)
+
+Using the `Rule Information` and `Details` columns capa explorer shows us that the suspect function matched `self delete via COMSPEC environment variable` because it contains capa rule matches for `create process`, `get COMSPEC environment variable`,
+and `query environment variable`, references to the strings `COMSPEC`, ` > nul`, and `/c del`, and calls to the Windows API functions `GetEnvironmentVariableA` and `ShellExecuteEx`.
+
+For more information on the FLARE team's open-source framework, capa, check out the overview in our first [blog](https://www.fireeye.com/blog/threat-research/2020/07/capa-automatically-identify-malware-capabilities.html).
+
+## Features
+
+![](../../../doc/img/ida_plugin_intro.gif)
+
+* Display capa results in an interactive tree view of rule matches and their locations in the current database
+* Search for keywords or phrases found in the `Rule Information`, `Address`, or `Details` columns
+* Display rule source content when a user hovers their cursor over a rule match
+* Double-click `Address` column to view associated feature in the IDA Disassembly view
+* Limit tree view results to the function currently displayed in the IDA Disassembly view; update results as a user navigates to different functions
+* Export results as formatted JSON by navigating to `File > Export results...`
+* Remember a user's capa rules directory for future runs; change capa rules directory by navigating to `Rules > Change rules directory...`
+* Automatically re-analyze database when user performs a program rebase
+* Automatically update results when IDA is used to rename a function
+* Select one or more checkboxes to highlight the associated addresses in the IDA Disassembly view
+* Right-click a function match to rename it; the new function name is propagated to the current IDA database
+* Right-click to copy a result by column or by row
+* Sort results by column
+* Reset tree view and IDA Disassembly view highlighting by clicking `Reset`
+
+## Getting Started
+
+### Requirements
+
+capa explorer supports the following IDA setups:
+
+* IDA Pro 7.4+ with Python 2.7 or Python 3.
+
+If you encounter issues with your specific setup, please open a new [Issue](https://github.com/fireeye/capa/issues).
+
+### Supported File Types
+
+capa explorer is limited to the file types supported by capa, which includes:
+
+* Windows 32-bit and 64-bit PE files
+* Windows 32-bit and 64-bit shellcode
+
+### Installation
+
+You can install capa explorer using the following steps:
+
+1. Install capa for the Python interpreter used by your IDA installation:
+    ```
+    $ pip install flare-capa
+    ```
+3. Download the [standard collection of capa rules](https://github.com/fireeye/capa-rules) (capa explorer needs capa rules to analyze a database)
+4. Copy [capa_explorer.py](https://raw.githubusercontent.com/fireeye/capa/master/capa/ida/plugin/capa_explorer.py) to your IDA plugins directory
+
+### Usage
+
+1. Run IDA and analyze a supported file type (select `Manual Load` and `Load Resources` for best results)
+2. Open capa explorer in IDA by navigating to `Edit > Plugins > FLARE capa explorer` or using the keyboard shortcut `Alt+F5`
+3. Click `Analyze`
+
+When running capa explorer for the first time you are prompted to select a file directory containing capa rules. The plugin conveniently
+remembers your selection for future runs; you can change this selection by navigating to `Rules > Change rules directory...`. We recommend 
+downloading and using the [standard collection of capa rules](https://github.com/fireeye/capa-rules) when getting started with the plugin.
+
+#### Tips
+
+* Start analysis by clicking `Analyze`
+* Reset the plugin user interface and remove highlighting from IDA disassembly view by clicking `Reset`
+* Change your capa rules directory by navigating to `Rules > Change rules directory...`
+* Hover your cursor over a rule match to view the source content of the rule
+* Double-click the `Address` column to navigate the IDA Disassembly view to the associated feature
+* Double-click a result in the `Rule Information` column to expand its children
+* Select a checkbox in the `Rule Information` column to highlight the address of the associated feature in the IDA Dissasembly view
+
+## Development
+
+Because capa explorer is packaged with capa you will need to install capa locally for development.
+
+You can install capa locally by following the steps outlined in `Method 3: Inspecting the capa source code` of the [capa 
+installation guide](https://github.com/fireeye/capa/blob/ida_plugin_documentation/doc/installation.md#method-3-inspecting-the-capa-source-code). Once installed, copy [capa_explorer.py](https://raw.githubusercontent.com/fireeye/capa/master/capa/ida/plugin/capa_explorer.py) 
+to your IDA plugins directory to run the plugin in IDA.
+
+### Components
+
+capa explorer consists of two main components:
+
+* An IDA [feature extractor](https://github.com/fireeye/capa/tree/master/capa/features/extractors/ida) built on top of IDA's binary analysis engine
+  * This component uses IDAPython to extract [capa features](https://github.com/fireeye/capa-rules/blob/master/doc/format.md#extracted-features) from the IDA database such as strings, 
+disassembly, and control flow; these extracted features are used by capa to find feature combinations that result in a rule match
+* An [interactive plugin](https://github.com/fireeye/capa/tree/master/capa/ida/plugin) for displaying and exploring capa rule matches
+  * This component integrates the IDA feature extractor and capa, providing an interactive user interface to dissect rule matches found by capa using features extracted by the IDA feature extractor

capa/ida/plugin/__init__.py

@@ -0,0 +1,117 @@
+# Copyright (C) 2020 FireEye, 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
+
+import idaapi
+import ida_kernwin
+
+from capa.ida.helpers import is_supported_file_type, is_supported_ida_version
+from capa.ida.plugin.form import CapaExplorerForm
+from capa.ida.plugin.icon import ICON
+
+logger = logging.getLogger(__name__)
+
+
+class CapaExplorerPlugin(idaapi.plugin_t):
+
+    # Mandatory definitions
+    PLUGIN_NAME = "FLARE capa explorer"
+    PLUGIN_VERSION = "1.0.0"
+    PLUGIN_AUTHORS = "michael.hunhoff@mandiant.com, william.ballenthin@mandiant.com, moritz.raabe@mandiant.com"
+
+    wanted_name = PLUGIN_NAME
+    wanted_hotkey = "ALT-F5"
+    comment = "IDA Pro plugin for the FLARE team's capa tool to identify capabilities in executable files."
+    website = "https://github.com/fireeye/capa"
+    help = "See https://github.com/fireeye/capa/blob/master/doc/usage.md"
+    version = ""
+    flags = 0
+
+    def __init__(self):
+        """initialize plugin"""
+        self.form = None
+
+    def init(self):
+        """called when IDA is loading the plugin"""
+        logging.basicConfig(level=logging.INFO)
+
+        # do not load plugin if IDA version/file type not supported
+        if not is_supported_ida_version():
+            return idaapi.PLUGIN_SKIP
+        if not is_supported_file_type():
+            return idaapi.PLUGIN_SKIP
+        return idaapi.PLUGIN_OK
+
+    def term(self):
+        """called when IDA is unloading the plugin"""
+        pass
+
+    def run(self, arg):
+        """called when IDA is running the plugin as a script"""
+        self.form = CapaExplorerForm(self.PLUGIN_NAME)
+        return True
+
+
+# set the capa plugin icon.
+#
+# TL;DR: temporarily install a UI hook set the icon.
+#
+# Long form:
+#
+# in the IDAPython `plugin_t` life cycle,
+#   - `init` decides if a plugin should be registered
+#   - `run` executes the main logic (shows the window)
+#   - `term` cleans this up
+#
+# we want to associate an icon with the plugin action - which is created by IDA.
+# however, this action is created by IDA *after* `init` is called.
+# so, we can't do this in `plugin_t.init`.
+# we also can't spawn a thread and do it after a delay,
+#  since `ida_kernwin.update_action_icon` must be called from the main thread.
+# so we need to register a callback that's invoked from the main thread after the plugin is registered.
+#
+# after a lot of guess-and-check, we can use `UI_Hooks.updated_actions` to
+#  receive notications after IDA has created an action for each plugin.
+# so, create this hook, wait for capa plugin to load, set the icon, and unhook.
+
+
+class OnUpdatedActionsHook(ida_kernwin.UI_Hooks):
+    """register a callback to be invoked each time the UI actions are updated"""
+
+    def __init__(self, cb):
+        super(OnUpdatedActionsHook, self).__init__()
+        self.cb = cb
+
+    def updated_actions(self):
+        if self.cb():
+            # uninstall the callback once its run successfully
+            self.unhook()
+
+
+def install_icon():
+    plugin_name = CapaExplorerPlugin.PLUGIN_NAME
+    action_name = "Edit/Plugins/" + plugin_name
+
+    if action_name not in ida_kernwin.get_registered_actions():
+        # keep the hook registered
+        return False
+
+    # resource leak here. need to call `ida_kernwin.free_custom_icon`?
+    # however, since we're not cycling this icon a lot, its probably ok.
+    # expect to leak exactly one icon per application load.
+    icon = ida_kernwin.load_custom_icon(data=ICON)
+
+    ida_kernwin.update_action_icon(action_name, icon)
+
+    # uninstall the hook
+    return True
+
+
+h = OnUpdatedActionsHook(install_icon)
+h.hook()

capa/ida/plugin/capa_explorer.py

@@ -0,0 +1,17 @@
+# Copyright (C) 2020 FireEye, 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 capa.ida.plugin import CapaExplorerPlugin
+
+
+def PLUGIN_ENTRY():
+    """mandatory entry point for IDAPython plugins
+
+    copy this script to your IDA plugins directory and start the plugin by navigating to Edit > Plugins in IDA Pro
+    """
+    return CapaExplorerPlugin()

capa/ida/plugin/form.py

@@ -0,0 +1,766 @@
+# Copyright (C) 2020 FireEye, 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 os
+import json
+import logging
+import collections
+
+import idaapi
+import ida_kernwin
+import ida_settings
+from PyQt5 import QtGui, QtCore, QtWidgets
+
+import capa.main
+import capa.rules
+import capa.ida.helpers
+import capa.render.utils as rutils
+import capa.features.extractors.ida
+from capa.ida.plugin.icon import QICON
+from capa.ida.plugin.view import CapaExplorerQtreeView
+from capa.ida.plugin.hooks import CapaExplorerIdaHooks
+from capa.ida.plugin.model import CapaExplorerDataModel
+from capa.ida.plugin.proxy import CapaExplorerRangeProxyModel, CapaExplorerSearchProxyModel
+
+logger = logging.getLogger(__name__)
+settings = ida_settings.IDASettings("capa")
+
+
+class UserCancelledError(Exception):
+    """throw exception when user cancels action"""
+
+    pass
+
+
+class CapaExplorerProgressIndicator(QtCore.QObject):
+    """implement progress signal, used during feature extraction"""
+
+    progress = QtCore.pyqtSignal(str)
+
+    def __init__(self):
+        """initialize signal object"""
+        super(CapaExplorerProgressIndicator, self).__init__()
+
+    def update(self, text):
+        """emit progress update
+
+        check if user cancelled action, raise exception for parent function to catch
+        """
+        if ida_kernwin.user_cancelled():
+            raise UserCancelledError("user cancelled")
+        self.progress.emit("extracting features from %s" % text)
+
+
+class CapaExplorerFeatureExtractor(capa.features.extractors.ida.IdaFeatureExtractor):
+    """subclass the IdaFeatureExtractor
+
+    track progress during feature extraction, also allow user to cancel feature extraction
+    """
+
+    def __init__(self):
+        super(CapaExplorerFeatureExtractor, self).__init__()
+        self.indicator = CapaExplorerProgressIndicator()
+
+    def extract_function_features(self, f):
+        self.indicator.update("function at 0x%X" % f.start_ea)
+        return super(CapaExplorerFeatureExtractor, self).extract_function_features(f)
+
+
+class CapaExplorerForm(idaapi.PluginForm):
+    """form element for plugin interface"""
+
+    def __init__(self, name):
+        """initialize form elements"""
+        super(CapaExplorerForm, self).__init__()
+
+        self.form_title = name
+        self.rule_path = ""
+        self.process_total = 0
+        self.process_count = 0
+
+        self.parent = None
+        self.ida_hooks = None
+        self.doc = None
+
+        # models
+        self.model_data = None
+        self.range_model_proxy = None
+        self.search_model_proxy = None
+
+        # UI controls
+        self.view_limit_results_by_function = None
+        self.view_search_bar = None
+        self.view_tree = None
+        self.view_attack = None
+        self.view_tabs = None
+        self.view_menu_bar = None
+        self.view_status_label = None
+        self.view_buttons = None
+        self.view_analyze_button = None
+        self.view_reset_button = None
+
+        self.Show()
+
+    def OnCreate(self, form):
+        """called when plugin form is created
+
+        load interface and install hooks but do not analyze database
+        """
+        self.parent = self.FormToPyQtWidget(form)
+        self.parent.setWindowIcon(QICON)
+        self.load_interface()
+        self.load_ida_hooks()
+
+    def Show(self):
+        """creates form if not already create, else brings plugin to front"""
+        return super(CapaExplorerForm, self).Show(
+            self.form_title,
+            options=(
+                idaapi.PluginForm.WOPN_TAB
+                | idaapi.PluginForm.WOPN_RESTORE
+                | idaapi.PluginForm.WCLS_CLOSE_LATER
+                | idaapi.PluginForm.WCLS_SAVE
+            ),
+        )
+
+    def OnClose(self, form):
+        """called when form is closed
+
+        ensure any plugin modifications (e.g. hooks and UI changes) are reset before the plugin is closed
+        """
+        self.unload_ida_hooks()
+        self.model_data.reset()
+
+    def load_interface(self):
+        """load user interface"""
+        # load models
+        self.model_data = CapaExplorerDataModel()
+
+        # model <- filter range <- filter search <- view
+
+        self.range_model_proxy = CapaExplorerRangeProxyModel()
+        self.range_model_proxy.setSourceModel(self.model_data)
+
+        self.search_model_proxy = CapaExplorerSearchProxyModel()
+        self.search_model_proxy.setSourceModel(self.range_model_proxy)
+
+        self.view_tree = CapaExplorerQtreeView(self.search_model_proxy, self.parent)
+        self.load_view_attack()
+
+        # load parent tab and children tab views
+        self.load_view_tabs()
+        self.load_view_checkbox_limit_by()
+        self.load_view_search_bar()
+        self.load_view_tree_tab()
+        self.load_view_attack_tab()
+        self.load_view_status_label()
+        self.load_view_buttons()
+
+        # load menu bar and sub menus
+        self.load_view_menu_bar()
+        self.load_file_menu()
+        self.load_rules_menu()
+
+        # load parent view
+        self.load_view_parent()
+
+        self.disable_controls()
+
+    def load_view_tabs(self):
+        """load tabs"""
+        tabs = QtWidgets.QTabWidget()
+        self.view_tabs = tabs
+
+    def load_view_menu_bar(self):
+        """load menu bar"""
+        bar = QtWidgets.QMenuBar()
+        self.view_menu_bar = bar
+
+    def load_view_attack(self):
+        """load MITRE ATT&CK table"""
+        table_headers = [
+            "ATT&CK Tactic",
+            "ATT&CK Technique ",
+        ]
+
+        table = QtWidgets.QTableWidget()
+
+        table.setColumnCount(len(table_headers))
+        table.verticalHeader().setVisible(False)
+        table.setSortingEnabled(False)
+        table.setEditTriggers(QtWidgets.QAbstractItemView.NoEditTriggers)
+        table.setFocusPolicy(QtCore.Qt.NoFocus)
+        table.setSelectionMode(QtWidgets.QAbstractItemView.NoSelection)
+        table.setHorizontalHeaderLabels(table_headers)
+        table.horizontalHeader().setDefaultAlignment(QtCore.Qt.AlignLeft)
+        table.setShowGrid(False)
+        table.setStyleSheet("QTableWidget::item { padding: 25px; }")
+
+        self.view_attack = table
+
+    def load_view_checkbox_limit_by(self):
+        """load limit results by function checkbox"""
+        check = QtWidgets.QCheckBox("Limit results to current function")
+        check.setChecked(False)
+        check.stateChanged.connect(self.slot_checkbox_limit_by_changed)
+
+        self.view_limit_results_by_function = check
+
+    def load_view_status_label(self):
+        """load status label"""
+        label = QtWidgets.QLabel()
+        label.setAlignment(QtCore.Qt.AlignLeft)
+        label.setText("Click Analyze to get started...")
+
+        self.view_status_label = label
+
+    def load_view_buttons(self):
+        """load the button controls"""
+        analyze_button = QtWidgets.QPushButton("Analyze")
+        analyze_button.setToolTip("Run capa analysis on IDB")
+        reset_button = QtWidgets.QPushButton("Reset")
+        reset_button.setToolTip("Reset capa explorer and IDA user interfaces")
+
+        analyze_button.clicked.connect(self.slot_analyze)
+        reset_button.clicked.connect(self.slot_reset)
+
+        layout = QtWidgets.QHBoxLayout()
+        layout.addWidget(analyze_button)
+        layout.addWidget(reset_button)
+        layout.addStretch(1)
+
+        self.view_analyze_button = analyze_button
+        self.view_reset_button = reset_button
+        self.view_buttons = layout
+
+    def load_view_search_bar(self):
+        """load the search bar control"""
+        line = QtWidgets.QLineEdit()
+        line.setPlaceholderText("search...")
+        line.textChanged.connect(self.slot_limit_results_to_search)
+
+        self.view_search_bar = line
+
+    def load_view_parent(self):
+        """load view parent"""
+        layout = QtWidgets.QVBoxLayout()
+
+        layout.addWidget(self.view_tabs)
+        layout.addWidget(self.view_status_label)
+        layout.addLayout(self.view_buttons)
+        layout.setMenuBar(self.view_menu_bar)
+
+        self.parent.setLayout(layout)
+
+    def load_view_tree_tab(self):
+        """load tree view tab"""
+        layout = QtWidgets.QVBoxLayout()
+        layout.addWidget(self.view_limit_results_by_function)
+        layout.addWidget(self.view_search_bar)
+        layout.addWidget(self.view_tree)
+
+        tab = QtWidgets.QWidget()
+        tab.setLayout(layout)
+
+        self.view_tabs.addTab(tab, "Tree View")
+
+    def load_view_attack_tab(self):
+        """load MITRE ATT&CK view tab"""
+        layout = QtWidgets.QVBoxLayout()
+        layout.addWidget(self.view_attack)
+
+        tab = QtWidgets.QWidget()
+        tab.setLayout(layout)
+
+        self.view_tabs.addTab(tab, "MITRE")
+
+    def load_file_menu(self):
+        """load file menu controls"""
+        actions = (("Export results...", "Export capa results as JSON file", self.slot_export_json),)
+        self.load_menu("File", actions)
+
+    def load_rules_menu(self):
+        """load rules menu controls"""
+        actions = (("Change rules directory...", "Select new rules directory", self.slot_change_rules_dir),)
+        self.load_menu("Rules", actions)
+
+    def load_menu(self, title, actions):
+        """load menu actions
+
+        @param title: menu name displayed in UI
+        @param actions: tuple of tuples containing action name, tooltip, and slot function
+        """
+        menu = self.view_menu_bar.addMenu(title)
+        for (name, _, slot) in actions:
+            action = QtWidgets.QAction(name, self.parent)
+            action.triggered.connect(slot)
+            menu.addAction(action)
+
+    def slot_export_json(self):
+        """export capa results as JSON file"""
+        if not self.doc:
+            idaapi.info("No capa results to export.")
+            return
+
+        path = idaapi.ask_file(True, "*.json", "Choose file")
+
+        # user cancelled, entered blank input, etc.
+        if not path:
+            return
+
+        # check file exists, ask to override
+        if os.path.exists(path) and 1 != idaapi.ask_yn(1, "The selected file already exists. Overwrite?"):
+            return
+
+        with open(path, "wb") as export_file:
+            export_file.write(
+                json.dumps(self.doc, sort_keys=True, cls=capa.render.CapaJsonObjectEncoder).encode("utf-8")
+            )
+
+    def load_ida_hooks(self):
+        """load IDA UI hooks"""
+        # map named action (defined in idagui.cfg) to Python function
+        action_hooks = {
+            "MakeName": self.ida_hook_rename,
+            "EditFunction": self.ida_hook_rename,
+            "RebaseProgram": self.ida_hook_rebase,
+        }
+
+        self.ida_hooks = CapaExplorerIdaHooks(self.ida_hook_screen_ea_changed, action_hooks)
+        self.ida_hooks.hook()
+
+    def unload_ida_hooks(self):
+        """unload IDA Pro UI hooks
+
+        must be called before plugin is completely destroyed
+        """
+        if self.ida_hooks:
+            self.ida_hooks.unhook()
+
+    def ida_hook_rename(self, meta, post=False):
+        """function hook for IDA "MakeName" and "EditFunction" actions
+
+        called twice, once before action and once after action completes
+
+        @param meta: dict of key/value pairs set when action first called (may be empty)
+        @param post: False if action first call, True if action second call
+        """
+        location = idaapi.get_screen_ea()
+        if not location or not capa.ida.helpers.is_func_start(location):
+            return
+
+        curr_name = idaapi.get_name(location)
+
+        if post:
+            # post action update data model w/ current name
+            self.model_data.update_function_name(meta.get("prev_name", ""), curr_name)
+        else:
+            # pre action so save current name for replacement later
+            meta["prev_name"] = curr_name
+
+    def ida_hook_screen_ea_changed(self, widget, new_ea, old_ea):
+        """function hook for IDA "screen ea changed" action
+
+        called twice, once before action and once after action completes. this hook is currently only relevant
+        for limiting results displayed in the UI
+
+        @param widget: IDA widget type
+        @param new_ea: destination ea
+        @param old_ea: source ea
+        """
+        if not self.view_limit_results_by_function.isChecked():
+            # ignore if limit checkbox not selected
+            return
+
+        if idaapi.get_widget_type(widget) != idaapi.BWN_DISASM:
+            # ignore views not the assembly view
+            return
+
+        if idaapi.get_func(new_ea) == idaapi.get_func(old_ea):
+            # user navigated same function - ignore
+            return
+
+        self.limit_results_to_function(idaapi.get_func(new_ea))
+        self.view_tree.reset_ui()
+
+    def ida_hook_rebase(self, meta, post=False):
+        """function hook for IDA "RebaseProgram" action
+
+        called twice, once before action and once after action completes
+
+        @param meta: dict of key/value pairs set when action first called (may be empty)
+        @param post: False if action first call, True if action second call
+        """
+        if post:
+            if idaapi.get_imagebase() != meta.get("prev_base", -1):
+                capa.ida.helpers.inform_user_ida_ui("Running capa analysis again after program rebase")
+                self.slot_analyze()
+        else:
+            meta["prev_base"] = idaapi.get_imagebase()
+            self.model_data.reset()
+
+    def load_capa_results(self):
+        """run capa analysis and render results in UI
+
+        note: this function must always return, exception or not, in order for plugin to safely close the IDA
+        wait box
+        """
+        # new analysis, new doc
+        self.doc = None
+        self.process_total = 0
+        self.process_count = 0
+
+        def update_wait_box(text):
+            """update the IDA wait box"""
+            ida_kernwin.replace_wait_box("capa explorer...%s" % text)
+
+        def slot_progress_feature_extraction(text):
+            """slot function to handle feature extraction progress updates"""
+            update_wait_box("%s (%d of %d)" % (text, self.process_count, self.process_total))
+            self.process_count += 1
+
+        extractor = CapaExplorerFeatureExtractor()
+        extractor.indicator.progress.connect(slot_progress_feature_extraction)
+
+        update_wait_box("calculating analysis")
+
+        try:
+            self.process_total += len(tuple(extractor.get_functions()))
+        except Exception as e:
+            logger.error("Failed to calculate analysis (error: %s).", e)
+            return False
+
+        if ida_kernwin.user_cancelled():
+            logger.info("User cancelled analysis.")
+            return False
+
+        update_wait_box("loading rules")
+
+        try:
+            # resolve rules directory - check self and settings first, then ask user
+            if not self.rule_path:
+                if "rule_path" in settings and os.path.exists(settings["rule_path"]):
+                    self.rule_path = settings["rule_path"]
+                else:
+                    idaapi.info("Please select a file directory containing capa rules.")
+                    rule_path = self.ask_user_directory()
+                    if not rule_path:
+                        logger.warning(
+                            "You must select a file directory containing capa rules before analysis can be run. The standard collection of capa rules can be downloaded from https://github.com/fireeye/capa-rules."
+                        )
+                        return False
+                    self.rule_path = rule_path
+                    settings.user["rule_path"] = rule_path
+        except Exception as e:
+            logger.error("Failed to load capa rules (error: %s).", e)
+            return False
+
+        if ida_kernwin.user_cancelled():
+            logger.info("User cancelled analysis.")
+            return False
+
+        rule_path = self.rule_path
+
+        try:
+            if not os.path.exists(rule_path):
+                raise IOError("rule path %s does not exist or cannot be accessed" % rule_path)
+
+            rule_paths = []
+            if os.path.isfile(rule_path):
+                rule_paths.append(rule_path)
+            elif os.path.isdir(rule_path):
+                for root, dirs, files in os.walk(rule_path):
+                    if ".github" in root:
+                        # the .github directory contains CI config in capa-rules
+                        # this includes some .yml files
+                        # these are not rules
+                        continue
+                    for file in files:
+                        if not file.endswith(".yml"):
+                            if not (file.endswith(".md") or file.endswith(".git") or file.endswith(".txt")):
+                                # expect to see readme.md, format.md, and maybe a .git directory
+                                # other things maybe are rules, but are mis-named.
+                                logger.warning("skipping non-.yml file: %s", file)
+                            continue
+                        rule_path = os.path.join(root, file)
+                        rule_paths.append(rule_path)
+
+            rules = []
+            total_paths = len(rule_paths)
+            for (i, rule_path) in enumerate(rule_paths):
+                update_wait_box("loading capa rules from %s (%d of %d)" % (self.rule_path, i + 1, total_paths))
+                if ida_kernwin.user_cancelled():
+                    raise UserCancelledError("user cancelled")
+                try:
+                    rule = capa.rules.Rule.from_yaml_file(rule_path)
+                except capa.rules.InvalidRule:
+                    raise
+                else:
+                    rule.meta["capa/path"] = rule_path
+                    if capa.main.is_nursery_rule_path(rule_path):
+                        rule.meta["capa/nursery"] = True
+                    rules.append(rule)
+
+            rule_count = len(rules)
+            rules = capa.rules.RuleSet(rules)
+        except UserCancelledError:
+            logger.info("User cancelled analysis.")
+            return False
+        except Exception as e:
+            capa.ida.helpers.inform_user_ida_ui("Failed to load capa rules from %s" % self.rule_path)
+            logger.error("Failed to load rules from %s (error: %s).", self.rule_path, e)
+            logger.error(
+                "Make sure your file directory contains properly formatted capa rules. You can download the standard collection of capa rules from https://github.com/fireeye/capa-rules."
+            )
+            self.rule_path = ""
+            settings.user.del_value("rule_path")
+            return False
+
+        if ida_kernwin.user_cancelled():
+            logger.info("User cancelled analysis.")
+            return False
+
+        update_wait_box("extracting features")
+
+        try:
+            meta = capa.ida.helpers.collect_metadata()
+            capabilities, counts = capa.main.find_capabilities(rules, extractor, disable_progress=True)
+            meta["analysis"].update(counts)
+        except UserCancelledError:
+            logger.info("User cancelled analysis.")
+            return False
+        except Exception as e:
+            logger.error("Failed to extract capabilities from database (error: %s)", e)
+            return False
+
+        update_wait_box("checking for file limitations")
+
+        try:
+            # support binary files specifically for x86/AMD64 shellcode
+            # warn user binary file is loaded but still allow capa to process it
+            # TODO: check specific architecture of binary files based on how user configured IDA processors
+            if idaapi.get_file_type_name() == "Binary file":
+                logger.warning("-" * 80)
+                logger.warning(" Input file appears to be a binary file.")
+                logger.warning(" ")
+                logger.warning(
+                    " capa currently only supports analyzing binary files containing x86/AMD64 shellcode with IDA."
+                )
+                logger.warning(
+                    " This means the results may be misleading or incomplete if the binary file loaded in IDA is not x86/AMD64."
+                )
+                logger.warning(
+                    " If you don't know the input file type, you can try using the `file` utility to guess it."
+                )
+                logger.warning("-" * 80)
+
+                capa.ida.helpers.inform_user_ida_ui("capa encountered file type warnings during analysis")
+
+            if capa.main.has_file_limitation(rules, capabilities, is_standalone=False):
+                capa.ida.helpers.inform_user_ida_ui("capa encountered file limitation warnings during analysis")
+        except Exception as e:
+            logger.error("Failed to check for file limitations (error: %s)", e)
+            return False
+
+        if ida_kernwin.user_cancelled():
+            logger.info("User cancelled analysis.")
+            return False
+
+        update_wait_box("rendering results")
+
+        try:
+            self.doc = capa.render.convert_capabilities_to_result_document(meta, rules, capabilities)
+            self.model_data.render_capa_doc(self.doc)
+            self.render_capa_doc_mitre_summary()
+            self.enable_controls()
+            self.set_view_status_label("capa rules directory: %s (%d rules)" % (self.rule_path, rule_count))
+        except Exception as e:
+            logger.error("Failed to render results (error: %s)", e)
+            return False
+
+        return True
+
+    def render_capa_doc_mitre_summary(self):
+        """render MITRE ATT&CK results"""
+        tactics = collections.defaultdict(set)
+
+        for rule in rutils.capability_rules(self.doc):
+            if not rule["meta"].get("att&ck"):
+                continue
+
+            for attack in rule["meta"]["att&ck"]:
+                tactic, _, rest = attack.partition("::")
+                if "::" in rest:
+                    technique, _, rest = rest.partition("::")
+                    subtechnique, _, id = rest.rpartition(" ")
+                    tactics[tactic].add((technique, subtechnique, id))
+                else:
+                    technique, _, id = rest.rpartition(" ")
+                    tactics[tactic].add((technique, id))
+
+        column_one = []
+        column_two = []
+
+        for (tactic, techniques) in sorted(tactics.items()):
+            column_one.append(tactic.upper())
+            # add extra space when more than one technique
+            column_one.extend(["" for i in range(len(techniques) - 1)])
+
+            for spec in sorted(techniques):
+                if len(spec) == 2:
+                    technique, id = spec
+                    column_two.append("%s %s" % (technique, id))
+                elif len(spec) == 3:
+                    technique, subtechnique, id = spec
+                    column_two.append("%s::%s %s" % (technique, subtechnique, id))
+                else:
+                    raise RuntimeError("unexpected ATT&CK spec format")
+
+        self.view_attack.setRowCount(max(len(column_one), len(column_two)))
+
+        for (row, value) in enumerate(column_one):
+            self.view_attack.setItem(row, 0, self.render_new_table_header_item(value))
+
+        for (row, value) in enumerate(column_two):
+            self.view_attack.setItem(row, 1, QtWidgets.QTableWidgetItem(value))
+
+        # resize columns to content
+        self.view_attack.resizeColumnsToContents()
+
+    def render_new_table_header_item(self, text):
+        """create new table header item with our style
+
+        @param text: header text to display
+        """
+        item = QtWidgets.QTableWidgetItem(text)
+        item.setForeground(QtGui.QColor(37, 147, 215))
+        font = QtGui.QFont()
+        font.setBold(True)
+        item.setFont(font)
+        return item
+
+    def reset_view_tree(self):
+        """reset tree view UI controls
+
+        called when user selects plugin reset from menu
+        """
+        self.view_limit_results_by_function.setChecked(False)
+        self.view_search_bar.setText("")
+        self.view_tree.reset_ui()
+
+    def slot_analyze(self):
+        """run capa analysis and reload UI controls
+
+        called when user selects plugin reload from menu
+        """
+        self.range_model_proxy.invalidate()
+        self.search_model_proxy.invalidate()
+        self.model_data.reset()
+        self.model_data.clear()
+        self.disable_controls()
+        self.set_view_status_label("Loading...")
+
+        ida_kernwin.show_wait_box("capa explorer")
+        success = self.load_capa_results()
+        ida_kernwin.hide_wait_box()
+
+        self.reset_view_tree()
+
+        if not success:
+            self.set_view_status_label("Click Analyze to get started...")
+            logger.info("Analysis failed.")
+        else:
+            logger.info("Analysis completed.")
+
+    def slot_reset(self, checked):
+        """reset UI elements
+
+        e.g. checkboxes and IDA highlighting
+        """
+        self.model_data.reset()
+        self.reset_view_tree()
+        logger.info("Reset completed.")
+
+    def slot_checkbox_limit_by_changed(self, state):
+        """slot activated if checkbox clicked
+
+        if checked, configure function filter if screen location is located in function, otherwise clear filter
+
+        @param state: checked state
+        """
+        if state == QtCore.Qt.Checked:
+            self.limit_results_to_function(idaapi.get_func(idaapi.get_screen_ea()))
+        else:
+            self.range_model_proxy.reset_address_range_filter()
+
+        self.view_tree.reset_ui()
+
+    def limit_results_to_function(self, f):
+        """add filter to limit results to current function
+
+        adds new address range filter to include function bounds, allowing basic blocks matched within a function
+        to be included in the results
+
+        @param f: (IDA func_t)
+        """
+        if f:
+            self.range_model_proxy.add_address_range_filter(f.start_ea, f.end_ea)
+        else:
+            # if function not exists don't display any results (assume address never -1)
+            self.range_model_proxy.add_address_range_filter(-1, -1)
+
+    def slot_limit_results_to_search(self, text):
+        """limit tree view results to search matches
+
+        reset view after filter to maintain level 1 expansion
+        """
+        self.search_model_proxy.set_query(text)
+        self.view_tree.reset_ui(should_sort=False)
+
+    def ask_user_directory(self):
+        """create Qt dialog to ask user for a directory"""
+        return str(
+            QtWidgets.QFileDialog.getExistingDirectory(
+                self.parent, "Please select a capa rules directory", self.rule_path
+            )
+        )
+
+    def slot_change_rules_dir(self):
+        """allow user to change rules directory
+
+        user selection stored in settings for future runs
+        """
+        rule_path = self.ask_user_directory()
+        if not rule_path:
+            logger.warning("No rule directory selected, nothing to do.")
+            return
+
+        self.rule_path = rule_path
+        settings.user["rule_path"] = rule_path
+
+        if 1 == idaapi.ask_yn(1, "Run analysis now?"):
+            self.slot_analyze()
+
+    def set_view_status_label(self, text):
+        """update status label control
+
+        @param text: updated text
+        """
+        self.view_status_label.setText(text)
+
+    def disable_controls(self):
+        """disable form controls"""
+        self.view_reset_button.setEnabled(False)
+        self.view_tabs.setTabEnabled(0, False)
+        self.view_tabs.setTabEnabled(1, False)
+
+    def enable_controls(self):
+        """enable form controls"""
+        self.view_reset_button.setEnabled(True)
+        self.view_tabs.setTabEnabled(0, True)
+        self.view_tabs.setTabEnabled(1, True)

capa/ida/plugin/hooks.py

@@ -0,0 +1,60 @@
+# Copyright (C) 2020 FireEye, 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 idaapi
+
+
+class CapaExplorerIdaHooks(idaapi.UI_Hooks):
+    def __init__(self, screen_ea_changed_hook, action_hooks):
+        """facilitate IDA UI hooks
+
+        @param screen_ea_changed_hook: function hook for IDA screen ea changed
+        @param action_hooks: dict of IDA action handles
+        """
+        super(CapaExplorerIdaHooks, self).__init__()
+
+        self.screen_ea_changed_hook = screen_ea_changed_hook
+        self.process_action_hooks = action_hooks
+        self.process_action_handle = None
+        self.process_action_meta = {}
+
+    def preprocess_action(self, name):
+        """called prior to action completed
+
+        @param name: name of action defined by idagui.cfg
+
+        @retval must be 0
+        """
+        self.process_action_handle = self.process_action_hooks.get(name, None)
+
+        if self.process_action_handle:
+            self.process_action_handle(self.process_action_meta)
+
+        # must return 0 for IDA
+        return 0
+
+    def postprocess_action(self):
+        """called after action completed"""
+        if not self.process_action_handle:
+            return
+
+        self.process_action_handle(self.process_action_meta, post=True)
+        self.reset()
+
+    def screen_ea_changed(self, curr_ea, prev_ea):
+        """called after screen location is changed
+
+        @param curr_ea: current location
+        @param prev_ea: prev location
+        """
+        self.screen_ea_changed_hook(idaapi.get_current_widget(), curr_ea, prev_ea)
+
+    def reset(self):
+        """reset internal state"""
+        self.process_action_handle = None
+        self.process_action_meta.clear()

capa/ida/plugin/item.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 sys
 import codecs
@@ -11,9 +17,9 @@ import capa.ida.helpers
 
 
 def info_to_name(display):
-    """ extract root value from display name
+    """extract root value from display name
 
-        e.g. function(my_function) => my_function
+    e.g. function(my_function) => my_function
     """
     try:
         return display.split("(")[1].rstrip(")")
@@ -22,20 +28,21 @@ def info_to_name(display):
 
 
 def location_to_hex(location):
-    """ convert location to hex for display """
+    """convert location to hex for display"""
     return "%08X" % location
 
 
 class CapaExplorerDataItem(object):
-    """ store data for CapaExplorerDataModel """
+    """store data for CapaExplorerDataModel"""
 
     def __init__(self, parent, data):
-        """ """
+        """initialize item"""
         self.pred = parent
         self._data = data
         self.children = []
         self._checked = False
 
+        # default state for item
         self.flags = (
             QtCore.Qt.ItemIsEnabled
             | QtCore.Qt.ItemIsSelectable
@@ -47,117 +54,146 @@ class CapaExplorerDataItem(object):
             self.pred.appendChild(self)
 
     def setIsEditable(self, isEditable=False):
-        """ modify item flags to be editable or not """
+        """modify item editable flags
+
+        @param isEditable: True, can edit, False cannot edit
+        """
         if isEditable:
             self.flags |= QtCore.Qt.ItemIsEditable
         else:
             self.flags &= ~QtCore.Qt.ItemIsEditable
 
     def setChecked(self, checked):
-        """ set item as checked """
+        """set item as checked
+
+        @param checked: True, item checked, False item not checked
+        """
         self._checked = checked
 
     def isChecked(self):
-        """ get item is checked """
+        """get item is checked"""
         return self._checked
 
     def appendChild(self, item):
-        """ add child item
+        """add a new child to specified item
 
-            @param item: CapaExplorerDataItem*
+        @param item: CapaExplorerDataItem
         """
         self.children.append(item)
 
     def child(self, row):
-        """ get child row
+        """get child row
 
-            @param row: TODO
+        @param row: row number
         """
         return self.children[row]
 
     def childCount(self):
-        """ get child count """
+        """get child count"""
         return len(self.children)
 
     def columnCount(self):
-        """ get column count """
+        """get column count"""
         return len(self._data)
 
     def data(self, column):
-        """ get data at column """
+        """get data at column
+
+        @param: column number
+        """
         try:
             return self._data[column]
         except IndexError:
             return None
 
     def parent(self):
-        """ get parent """
+        """get parent"""
         return self.pred
 
     def row(self):
-        """ get row location """
+        """get row location"""
         if self.pred:
             return self.pred.children.index(self)
         return 0
 
     def setData(self, column, value):
-        """ set data in column """
+        """set data in column
+
+        @param column: column number
+        @value: value to set (assume str)
+        """
         self._data[column] = value
 
     def children(self):
-        """ yield children """
+        """yield children"""
         for child in self.children:
             yield child
 
     def removeChildren(self):
-        """ remove children from node """
+        """remove children"""
         del self.children[:]
 
     def __str__(self):
-        """ get string representation of columns """
+        """get string representation of columns
+
+        used for copy-n-paste operations
+        """
         return " ".join([data for data in self._data if data])
 
     @property
     def info(self):
-        """ return data stored in information column """
+        """return data stored in information column"""
         return self._data[0]
 
     @property
     def location(self):
-        """ return data stored in location column """
+        """return data stored in location column"""
         try:
+            # address stored as str, convert to int before return
             return int(self._data[1], 16)
         except ValueError:
             return None
 
     @property
     def details(self):
-        """ return data stored in details column """
+        """return data stored in details column"""
         return self._data[2]
 
 
 class CapaExplorerRuleItem(CapaExplorerDataItem):
-    """ store data relevant to capa function result """
+    """store data for rule result"""
 
     fmt = "%s (%d matches)"
 
-    def __init__(self, parent, display, count, source):
-        """ """
-        display = self.fmt % (display, count) if count > 1 else display
-        super(CapaExplorerRuleItem, self).__init__(parent, [display, "", ""])
+    def __init__(self, parent, name, namespace, count, source):
+        """initialize item
+
+        @param parent: parent node
+        @param name: rule name
+        @param namespace: rule namespace
+        @param count: number of match for this rule
+        @param source: rule source (tooltip)
+        """
+        display = self.fmt % (name, count) if count > 1 else name
+        super(CapaExplorerRuleItem, self).__init__(parent, [display, "", namespace])
         self._source = source
 
     @property
     def source(self):
-        """ return rule contents for display """
+        """return rule source to display (tooltip)"""
         return self._source
 
 
 class CapaExplorerRuleMatchItem(CapaExplorerDataItem):
-    """ store data relevant to capa function match result """
+    """store data for rule match"""
 
     def __init__(self, parent, display, source=""):
-        """ """
+        """initialize item
+
+        @param parent: parent node
+        @param display: text to display in UI
+        @param source: rule match source to display (tooltip)
+        """
         super(CapaExplorerRuleMatchItem, self).__init__(parent, [display, "", ""])
         self._source = source
 
@@ -168,82 +204,125 @@ class CapaExplorerRuleMatchItem(CapaExplorerDataItem):
 
 
 class CapaExplorerFunctionItem(CapaExplorerDataItem):
-    """ store data relevant to capa function result """
+    """store data for function match"""
 
     fmt = "function(%s)"
 
     def __init__(self, parent, location):
-        """ """
+        """initialize item
+
+        @param parent: parent node
+        @param location: virtual address of function as seen by IDA
+        """
         super(CapaExplorerFunctionItem, self).__init__(
             parent, [self.fmt % idaapi.get_name(location), location_to_hex(location), ""]
         )
 
     @property
     def info(self):
-        """ """
+        """return function name"""
         info = super(CapaExplorerFunctionItem, self).info
         display = info_to_name(info)
         return display if display else info
 
     @info.setter
     def info(self, display):
-        """ """
+        """set function name
+
+        called when user changes function name in plugin UI
+
+        @param display: new function name to display
+        """
         self._data[0] = self.fmt % display
 
 
 class CapaExplorerSubscopeItem(CapaExplorerDataItem):
-    """ store data relevant to subscope """
+    """store data for subscope match"""
 
     fmt = "subscope(%s)"
 
     def __init__(self, parent, scope):
-        """ """
+        """initialize item
+
+        @param parent: parent node
+        @param scope: subscope name
+        """
         super(CapaExplorerSubscopeItem, self).__init__(parent, [self.fmt % scope, "", ""])
 
 
 class CapaExplorerBlockItem(CapaExplorerDataItem):
-    """ store data relevant to capa basic block result """
+    """store data for basic block match"""
 
     fmt = "basic block(loc_%08X)"
 
     def __init__(self, parent, location):
-        """ """
+        """initialize item
+
+        @param parent: parent node
+        @param location: virtual address of basic block as seen by IDA
+        """
         super(CapaExplorerBlockItem, self).__init__(parent, [self.fmt % location, location_to_hex(location), ""])
 
 
 class CapaExplorerDefaultItem(CapaExplorerDataItem):
-    """ store data relevant to capa default result """
+    """store data for default match e.g. statement (and, or)"""
 
     def __init__(self, parent, display, details="", location=None):
-        """ """
+        """initialize item
+
+        @param parent: parent node
+        @param display: text to display in UI
+        @param details: text to display in details section of UI
+        @param location: virtual address as seen by IDA
+        """
         location = location_to_hex(location) if location else ""
         super(CapaExplorerDefaultItem, self).__init__(parent, [display, location, details])
 
 
 class CapaExplorerFeatureItem(CapaExplorerDataItem):
-    """ store data relevant to capa feature result """
+    """store data for feature match"""
 
     def __init__(self, parent, display, location="", details=""):
-        """ """
+        """initialize item
+
+        @param parent: parent node
+        @param display: text to display in UI
+        @param details: text to display in details section of UI
+        @param location: virtual address as seen by IDA
+        """
         location = location_to_hex(location) if location else ""
         super(CapaExplorerFeatureItem, self).__init__(parent, [display, location, details])
 
 
 class CapaExplorerInstructionViewItem(CapaExplorerFeatureItem):
-    """ store data relevant to an instruction preview """
+    """store data for instruction match"""
 
     def __init__(self, parent, display, location):
-        """ """
+        """initialize item
+
+        details section shows disassembly view for match
+
+        @param parent: parent node
+        @param display: text to display in UI
+        @param location: virtual address as seen by IDA
+        """
         details = capa.ida.helpers.get_disasm_line(location)
         super(CapaExplorerInstructionViewItem, self).__init__(parent, display, location=location, details=details)
         self.ida_highlight = idc.get_color(location, idc.CIC_ITEM)
 
 
 class CapaExplorerByteViewItem(CapaExplorerFeatureItem):
-    """ store data relevant to byte preview """
+    """store data for byte match"""
 
     def __init__(self, parent, display, location):
-        """ """
+        """initialize item
+
+        details section shows byte preview for match
+
+        @param parent: parent node
+        @param display: text to display in UI
+        @param location: virtual address as seen by IDA
+        """
         byte_snap = idaapi.get_bytes(location, 32)
 
         if byte_snap:
@@ -260,9 +339,14 @@ class CapaExplorerByteViewItem(CapaExplorerFeatureItem):
 
 
 class CapaExplorerStringViewItem(CapaExplorerFeatureItem):
-    """ store data relevant to string preview """
+    """store data for string match"""
 
-    def __init__(self, parent, display, location):
-        """ """
-        super(CapaExplorerStringViewItem, self).__init__(parent, display, location=location)
+    def __init__(self, parent, display, location, value):
+        """initialize item
+
+        @param parent: parent node
+        @param display: text to display in UI
+        @param location: virtual address as seen by IDA
+        """
+        super(CapaExplorerStringViewItem, self).__init__(parent, display, location=location, details=value)
         self.ida_highlight = idc.get_color(location, idc.CIC_ITEM)

capa/ida/plugin/model.py

@@ -1,15 +1,20 @@
 # Copyright (C) 2020 FireEye, 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 collections import deque
 
 import idc
-import six
 import idaapi
-from PyQt5 import Qt, QtGui, QtCore
+from PyQt5 import QtGui, QtCore
 
 import capa.ida.helpers
 import capa.render.utils as rutils
-from capa.ida.explorer.item import (
+from capa.ida.plugin.item import (
     CapaExplorerDataItem,
     CapaExplorerRuleItem,
     CapaExplorerBlockItem,
@@ -24,11 +29,11 @@ from capa.ida.explorer.item import (
 )
 
 # default highlight color used in IDA window
-DEFAULT_HIGHLIGHT = 0xD096FF
+DEFAULT_HIGHLIGHT = 0xE6C700
 
 
 class CapaExplorerDataModel(QtCore.QAbstractItemModel):
-    """ """
+    """model for displaying hierarchical results return by capa"""
 
     COLUMN_INDEX_RULE_INFORMATION = 0
     COLUMN_INDEX_VIRTUAL_ADDRESS = 1
@@ -37,14 +42,16 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
     COLUMN_COUNT = 3
 
     def __init__(self, parent=None):
-        """ """
+        """initialize model"""
         super(CapaExplorerDataModel, self).__init__(parent)
+        # root node does not have parent, contains header columns
         self.root_node = CapaExplorerDataItem(None, ["Rule Information", "Address", "Details"])
 
     def reset(self):
-        """ """
-        # reset checkboxes and color highlights
-        # TODO: make less hacky
+        """reset UI elements (e.g. checkboxes, IDA color highlights)
+
+        called when view wants to reset UI display
+        """
         for idx in range(self.root_node.childCount()):
             root_index = self.index(idx, 0, QtCore.QModelIndex())
             for model_index in self.iterateChildrenIndexFromRootIndex(root_index, ignore_root=False):
@@ -53,17 +60,20 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
                 self.dataChanged.emit(model_index, model_index)
 
     def clear(self):
-        """ """
+        """clear model data
+
+        called when view wants to clear UI display
+        """
         self.beginResetModel()
         self.root_node.removeChildren()
         self.endResetModel()
 
     def columnCount(self, model_index):
-        """ get the number of columns for the children of the given parent
+        """return number of columns for the children of the given parent
 
-            @param model_index: QModelIndex*
+        @param model_index: QModelIndex
 
-            @retval column count
+        @retval column count
         """
         if model_index.isValid():
             return model_index.internalPointer().columnCount()
@@ -71,12 +81,14 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
             return self.root_node.columnCount()
 
     def data(self, model_index, role):
-        """ get data stored under the given role for the item referred to by the index
+        """return data stored at given index by display role
 
-            @param model_index: QModelIndex*
-            @param role: QtCore.Qt.*
+        this function is used to control UI elements (e.g. text font, color, etc.) based on column, item type, etc.
 
-            @retval data to be displayed
+        @param model_index: QModelIndex
+        @param role: QtCore.Qt.*
+
+        @retval data to be displayed
         """
         if not model_index.isValid():
             return None
@@ -125,14 +137,14 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
             )
             and column == CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION
         ):
-            # set bold font for top-level rules
+            # set bold font for important items
             font = QtGui.QFont()
             font.setBold(True)
             return font
 
         if role == QtCore.Qt.ForegroundRole and column == CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS:
             # set color for virtual address column
-            return QtGui.QColor(88, 139, 174)
+            return QtGui.QColor(37, 147, 215)
 
         if (
             role == QtCore.Qt.ForegroundRole
@@ -145,11 +157,11 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
         return None
 
     def flags(self, model_index):
-        """ get item flags for given index
+        """return item flags for given index
 
-            @param model_index: QModelIndex*
+        @param model_index: QModelIndex
 
-            @retval QtCore.Qt.ItemFlags
+        @retval QtCore.Qt.ItemFlags
         """
         if not model_index.isValid():
             return QtCore.Qt.NoItemFlags
@@ -157,13 +169,13 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
         return model_index.internalPointer().flags
 
     def headerData(self, section, orientation, role):
-        """ get data for the given role and section in the header with the specified orientation
+        """return data for the given role and section in the header with the specified orientation
 
-            @param section: int
-            @param orientation: QtCore.Qt.Orientation
-            @param role: QtCore.Qt.DisplayRole
+        @param section: int
+        @param orientation: QtCore.Qt.Orientation
+        @param role: QtCore.Qt.DisplayRole
 
-            @retval header data list()
+        @retval header data
         """
         if orientation == QtCore.Qt.Horizontal and role == QtCore.Qt.DisplayRole:
             return self.root_node.data(section)
@@ -171,13 +183,13 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
         return None
 
     def index(self, row, column, parent):
-        """ get index of the item in the model specified by the given row, column and parent index
+        """return index of the item by row, column, and parent index
 
-            @param row: int
-            @param column: int
-            @param parent: QModelIndex*
+        @param row: item row
+        @param column: item column
+        @param parent: QModelIndex of parent
 
-            @retval QModelIndex*
+        @retval QModelIndex of item
         """
         if not self.hasIndex(row, column, parent):
             return QtCore.QModelIndex()
@@ -195,13 +207,13 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
             return QtCore.QModelIndex()
 
     def parent(self, model_index):
-        """ get parent of the model item with the given index
+        """return parent index by child index
 
-            if the item has no parent, an invalid QModelIndex* is returned
+        if the item has no parent, an invalid QModelIndex is returned
 
-            @param model_index: QModelIndex*
+        @param model_index: QModelIndex of child
 
-            @retval QModelIndex*
+        @retval QModelIndex of parent
         """
         if not model_index.isValid():
             return QtCore.QModelIndex()
@@ -215,12 +227,12 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
         return self.createIndex(parent.row(), 0, parent)
 
     def iterateChildrenIndexFromRootIndex(self, model_index, ignore_root=True):
-        """ depth-first traversal of child nodes
+        """depth-first traversal of child nodes
 
-            @param model_index: QModelIndex*
-            @param ignore_root: if set, do not return root index
+        @param model_index: QModelIndex of starting item
+        @param ignore_root: True, do not yield root index, False yield root index
 
-            @retval yield QModelIndex*
+        @retval yield QModelIndex
         """
         visited = set()
         stack = deque((model_index,))
@@ -242,10 +254,10 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
                     stack.append(child_index.child(idx, 0))
 
     def reset_ida_highlighting(self, item, checked):
-        """ reset IDA highlight for an item
+        """reset IDA highlight for item
 
-            @param item: capa explorer item
-            @param checked: indicates item is or not checked
+        @param item: CapaExplorerDataItem
+        @param checked: True, item checked, False item not checked
         """
         if not isinstance(
             item, (CapaExplorerStringViewItem, CapaExplorerInstructionViewItem, CapaExplorerByteViewItem)
@@ -269,13 +281,11 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
                 idc.set_color(item.location, idc.CIC_ITEM, item.ida_highlight)
 
     def setData(self, model_index, value, role):
-        """ set the role data for the item at index to value
+        """set data at index by role
 
-            @param model_index: QModelIndex*
-            @param value: QVariant*
-            @param role: QtCore.Qt.EditRole
-
-            @retval True/False
+        @param model_index: QModelIndex of item
+        @param value: value to set
+        @param role: QtCore.Qt.EditRole
         """
         if not model_index.isValid():
             return False
@@ -310,14 +320,13 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
         return False
 
     def rowCount(self, model_index):
-        """ get the number of rows under the given parent
+        """return number of rows under item by index
 
-            when the parent is valid it means that is returning the number of
-            children of parent
+        when the parent is valid it means that is returning the number of children of parent
 
-            @param model_index: QModelIndex*
+        @param model_index: QModelIndex
 
-            @retval row count
+        @retval row count
         """
         if model_index.column() > 0:
             return 0
@@ -330,24 +339,26 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
         return item.childCount()
 
     def render_capa_doc_statement_node(self, parent, statement, locations, doc):
-        """ render capa statement read from doc
+        """render capa statement read from doc
 
-            @param parent: parent to which new child is assigned
-            @param statement: statement read from doc
-            @param locations: locations of children (applies to range only?)
-            @param doc: capa result doc
-
-            "statement": {
-                "type": "or"
-            },
+        @param parent: parent to which new child is assigned
+        @param statement: statement read from doc
+        @param locations: locations of children (applies to range only?)
+        @param doc: result doc
         """
         if statement["type"] in ("and", "or", "optional"):
-            return CapaExplorerDefaultItem(parent, statement["type"])
+            display = statement["type"]
+            if statement.get("description"):
+                display += " (%s)" % statement["description"]
+            return CapaExplorerDefaultItem(parent, display)
         elif statement["type"] == "not":
             # TODO: do we display 'not'
             pass
         elif statement["type"] == "some":
-            return CapaExplorerDefaultItem(parent, statement["count"] + " or more")
+            display = "%d or more" % statement["count"]
+            if statement.get("description"):
+                display += " (%s)" % statement["description"]
+            return CapaExplorerDefaultItem(parent, display)
         elif statement["type"] == "range":
             # `range` is a weird node, its almost a hybrid of statement + feature.
             # it is a specific feature repeated multiple times.
@@ -364,6 +375,9 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
             else:
                 display += "between %d and %d" % (statement["min"], statement["max"])
 
+            if statement.get("description"):
+                display += " (%s)" % statement["description"]
+
             parent2 = CapaExplorerFeatureItem(parent, display=display)
 
             for location in locations:
@@ -372,33 +386,19 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
 
             return parent2
         elif statement["type"] == "subscope":
-            return CapaExplorerSubscopeItem(parent, statement[statement["type"]])
+            display = statement[statement["type"]]
+            if statement.get("description"):
+                display += " (%s)" % statement["description"]
+            return CapaExplorerSubscopeItem(parent, display)
         else:
             raise RuntimeError("unexpected match statement type: " + str(statement))
 
     def render_capa_doc_match(self, parent, match, doc):
-        """ render capa match read from doc
+        """render capa match read from doc
 
-            @param parent: parent node to which new child is assigned
-            @param match: match read from doc
-            @param doc: capa result doc
-
-            "matches": {
-                "0": {
-                    "children": [],
-                    "locations": [
-                        4317184
-                    ],
-                    "node": {
-                        "feature": {
-                            "section": ".rsrc",
-                            "type": "section"
-                        },
-                        "type": "feature"
-                    },
-                    "success": true
-                }
-            },
+        @param parent: parent node to which new child is assigned
+        @param match: match read from doc
+        @param doc: result doc
         """
         if not match["success"]:
             # TODO: display failed branches at some point? Help with debugging rules?
@@ -425,15 +425,19 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
             self.render_capa_doc_match(parent2, child, doc)
 
     def render_capa_doc(self, doc):
-        """ render capa features specified in doc
+        """render capa features specified in doc
 
-            @param doc: capa result doc
+        @param doc: capa result doc
         """
         # inform model that changes are about to occur
         self.beginResetModel()
 
         for rule in rutils.capability_rules(doc):
-            parent = CapaExplorerRuleItem(self.root_node, rule["meta"]["name"], len(rule["matches"]), rule["source"])
+            rule_name = rule["meta"]["name"]
+            rule_namespace = rule["meta"].get("namespace")
+            parent = CapaExplorerRuleItem(
+                self.root_node, rule_name, rule_namespace, len(rule["matches"]), rule["source"]
+            )
 
             for (location, match) in doc["rules"][rule["meta"]["name"]]["matches"].items():
                 if rule["meta"]["scope"] == capa.rules.FILE_SCOPE:
@@ -451,18 +455,9 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
         self.endResetModel()
 
     def capa_doc_feature_to_display(self, feature):
-        """ convert capa doc feature type string to display string for ui
+        """convert capa doc feature type string to display string for ui
 
-            @param feature: capa feature read from doc
-
-            Example:
-                "feature": {
-                    "bytes": "01 14 02 00 00 00 00 00 C0 00 00 00 00 00 00 46",
-                    "description": "CLSID_ShellLink",
-                    "type": "bytes"
-                }
-
-                bytes(01 14 02 00 00 00 00 00 C0 00 00 00 00 00 00 46 = CLSID_ShellLink)
+        @param feature: capa feature read from doc
         """
         if feature[feature["type"]]:
             if feature.get("description", ""):
@@ -473,25 +468,24 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
             return "%s" % feature["type"]
 
     def render_capa_doc_feature_node(self, parent, feature, locations, doc):
-        """ process capa doc feature node
+        """process capa doc feature node
 
-            @param parent: parent node to which child is assigned
-            @param feature: capa doc feature node
-            @param locations: locations identified for feature
-            @param doc: capa doc
-
-            Example:
-              "feature": {
-                "description": "FILE_WRITE_DATA",
-                "number": "0x2",
-                "type": "number"
-              }
+        @param parent: parent node to which child is assigned
+        @param feature: capa doc feature node
+        @param locations: locations identified for feature
+        @param doc: capa doc
         """
         display = self.capa_doc_feature_to_display(feature)
 
         if len(locations) == 1:
             # only one location for feature so no need to nest children
-            parent2 = self.render_capa_doc_feature(parent, feature, next(iter(locations)), doc, display=display,)
+            parent2 = self.render_capa_doc_feature(
+                parent,
+                feature,
+                next(iter(locations)),
+                doc,
+                display=display,
+            )
         else:
             # feature has multiple children, nest  under one parent feature node
             parent2 = CapaExplorerFeatureItem(parent, display)
@@ -502,27 +496,20 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
         return parent2
 
     def render_capa_doc_feature(self, parent, feature, location, doc, display="-"):
-        """ render capa feature read from doc
+        """render capa feature read from doc
 
-            @param parent: parent node to which new child is assigned
-            @param feature: feature read from doc
-            @param doc: capa feature doc
-            @param location: address of feature
-            @param display: text to display in plugin ui
-
-            Example:
-              "feature": {
-                "description": "FILE_WRITE_DATA",
-                "number": "0x2",
-                "type": "number"
-              }
+        @param parent: parent node to which new child is assigned
+        @param feature: feature read from doc
+        @param doc: capa feature doc
+        @param location: address of feature
+        @param display: text to display in plugin UI
         """
         # special handling for characteristic pending type
         if feature["type"] == "characteristic":
             if feature[feature["type"]] in ("embedded pe",):
                 return CapaExplorerByteViewItem(parent, display, location)
 
-            if feature[feature["type"]] in ("loop", "recursive call", "tight loop", "switch"):
+            if feature[feature["type"]] in ("loop", "recursive call", "tight loop"):
                 return CapaExplorerFeatureItem(parent, display=display)
 
             # default to instruction view for all other characteristics
@@ -535,12 +522,22 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
             )
 
         if feature["type"] == "regex":
-            return CapaExplorerFeatureItem(parent, display, location, details=feature["match"])
+            return CapaExplorerStringViewItem(parent, display, location, feature["match"])
 
         if feature["type"] == "basicblock":
             return CapaExplorerBlockItem(parent, location)
 
-        if feature["type"] in ("bytes", "api", "mnemonic", "number", "offset"):
+        if feature["type"] in (
+            "bytes",
+            "api",
+            "mnemonic",
+            "number",
+            "offset",
+            "number/x32",
+            "number/x64",
+            "offset/x32",
+            "offset/x64",
+        ):
             # display instruction preview
             return CapaExplorerInstructionViewItem(parent, display, location)
 
@@ -550,7 +547,7 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
 
         if feature["type"] in ("string",):
             # display string preview
-            return CapaExplorerStringViewItem(parent, display, location)
+            return CapaExplorerStringViewItem(parent, display, location, feature[feature["type"]])
 
         if feature["type"] in ("import", "export"):
             # display no preview
@@ -559,10 +556,12 @@ class CapaExplorerDataModel(QtCore.QAbstractItemModel):
         raise RuntimeError("unexpected feature type: " + str(feature["type"]))
 
     def update_function_name(self, old_name, new_name):
-        """ update all instances of old function name with new function name
+        """update all instances of old function name with new function name
 
-            @param old_name: previous function name
-            @param new_name: new function name
+        called when user updates function name using plugin UI
+
+        @param old_name: old function name
+        @param new_name: new function name
         """
         # create empty root index for search
         root_index = self.index(0, 0, QtCore.QModelIndex())

capa/ida/plugin/proxy.py

@@ -0,0 +1,226 @@
+# Copyright (C) 2020 FireEye, 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 six
+from PyQt5 import QtCore
+from PyQt5.QtCore import Qt
+
+from capa.ida.plugin.model import CapaExplorerDataModel
+
+
+class CapaExplorerRangeProxyModel(QtCore.QSortFilterProxyModel):
+    """filter results based on virtual address range as seen by IDA
+
+    implements filtering for "limit results by current function" checkbox in plugin UI
+
+    minimum and maximum virtual addresses are used to filter results to a specific address range. this allows
+    basic blocks to be included when limiting results to a specific function
+    """
+
+    def __init__(self, parent=None):
+        """initialize proxy filter"""
+        super(CapaExplorerRangeProxyModel, self).__init__(parent)
+        self.min_ea = None
+        self.max_ea = None
+
+    def lessThan(self, left, right):
+        """return True if left item is less than right item, else False
+
+        @param left: QModelIndex of left
+        @param right: QModelIndex of right
+        """
+        ldata = left.internalPointer().data(left.column())
+        rdata = right.internalPointer().data(right.column())
+
+        if (
+            ldata
+            and rdata
+            and left.column() == CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS
+            and left.column() == right.column()
+        ):
+            # convert virtual address before compare
+            return int(ldata, 16) < int(rdata, 16)
+        else:
+            # compare as lowercase
+            return ldata.lower() < rdata.lower()
+
+    def filterAcceptsRow(self, row, parent):
+        """return true if the item in the row indicated by the given row and parent should be included in the model;
+        otherwise return false
+
+        @param row: row number
+        @param parent: QModelIndex of parent
+        """
+        if self.filter_accepts_row_self(row, parent):
+            return True
+
+        alpha = parent
+        while alpha.isValid():
+            if self.filter_accepts_row_self(alpha.row(), alpha.parent()):
+                return True
+            alpha = alpha.parent()
+
+        if self.index_has_accepted_children(row, parent):
+            return True
+
+        return False
+
+    def index_has_accepted_children(self, row, parent):
+        """return True if parent has one or more children that match filter, else False
+
+        @param row: row number
+        @param parent: QModelIndex of parent
+        """
+        model_index = self.sourceModel().index(row, 0, parent)
+
+        if model_index.isValid():
+            for idx in range(self.sourceModel().rowCount(model_index)):
+                if self.filter_accepts_row_self(idx, model_index):
+                    return True
+                if self.index_has_accepted_children(idx, model_index):
+                    return True
+
+        return False
+
+    def filter_accepts_row_self(self, row, parent):
+        """return True if filter accepts row, else False
+
+        @param row: row number
+        @param parent: QModelIndex of parent
+        """
+        # filter not set
+        if self.min_ea is None and self.max_ea is None:
+            return True
+
+        index = self.sourceModel().index(row, 0, parent)
+        data = index.internalPointer().data(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS)
+
+        # virtual address may be empty
+        if not data:
+            return False
+
+        # convert virtual address str to int
+        ea = int(data, 16)
+
+        if self.min_ea <= ea and ea < self.max_ea:
+            return True
+
+        return False
+
+    def add_address_range_filter(self, min_ea, max_ea):
+        """add new address range filter
+
+        called when user checks "limit results by current function" in plugin UI
+
+        @param min_ea: minimum virtual address as seen by IDA
+        @param max_ea: maximum virtual address as seen by IDA
+        """
+        self.min_ea = min_ea
+        self.max_ea = max_ea
+
+        self.setFilterKeyColumn(CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS)
+        self.invalidateFilter()
+
+    def reset_address_range_filter(self):
+        """remove address range filter (accept all results)
+
+        called when user un-checks "limit results by current function" in plugin UI
+        """
+        self.min_ea = None
+        self.max_ea = None
+        self.invalidateFilter()
+
+
+class CapaExplorerSearchProxyModel(QtCore.QSortFilterProxyModel):
+    """A SortFilterProxyModel that accepts rows with a substring match for a configurable query.
+
+    Looks for matches in the text of all rows.
+    Displays the entire tree row if any of the tree branches,
+     that is, you can filter by rule name, or also
+     filter by "characteristic(nzxor)" to filter matches with some feature.
+    """
+
+    def __init__(self, parent=None):
+        """ """
+        super(CapaExplorerSearchProxyModel, self).__init__(parent)
+        self.query = ""
+        self.setFilterKeyColumn(-1)  # all columns
+
+    def filterAcceptsRow(self, row, parent):
+        """true if the item in the row indicated by the given row and parent
+        should be included in the model; otherwise returns false
+
+        @param row: int
+        @param parent: QModelIndex*
+
+        @retval True/False
+        """
+        # this row matches, accept it
+        if self.filter_accepts_row_self(row, parent):
+            return True
+
+        # the parent of this row matches, accept it
+        alpha = parent
+        while alpha.isValid():
+            if self.filter_accepts_row_self(alpha.row(), alpha.parent()):
+                return True
+            alpha = alpha.parent()
+
+        # this row is a parent, and a child matches, accept it
+        if self.index_has_accepted_children(row, parent):
+            return True
+
+        return False
+
+    def index_has_accepted_children(self, row, parent):
+        """returns True if the given row or its children should be accepted"""
+        source_model = self.sourceModel()
+        model_index = source_model.index(row, 0, parent)
+
+        if model_index.isValid():
+            for idx in range(source_model.rowCount(model_index)):
+                if self.filter_accepts_row_self(idx, model_index):
+                    return True
+                if self.index_has_accepted_children(idx, model_index):
+                    return True
+
+        return False
+
+    def filter_accepts_row_self(self, row, parent):
+        """returns True if the given row should be accepted"""
+        if self.query == "":
+            return True
+
+        source_model = self.sourceModel()
+
+        for column in (
+            CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION,
+            CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS,
+            CapaExplorerDataModel.COLUMN_INDEX_DETAILS,
+        ):
+            index = source_model.index(row, column, parent)
+            data = source_model.data(index, Qt.DisplayRole)
+
+            if not data:
+                continue
+
+            if not isinstance(data, six.string_types):
+                # sanity check: should already be a string, but double check
+                continue
+
+            # case in-sensitive matching
+            if self.query.lower() in data.lower():
+                return True
+
+        return False
+
+    def set_query(self, query):
+        self.query = query
+        self.invalidateFilter()
+
+    def reset_query(self):
+        self.set_query("")

capa/ida/plugin/view.py

@@ -0,0 +1,291 @@
+# Copyright (C) 2020 FireEye, 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 idc
+from PyQt5 import QtCore, QtWidgets
+
+from capa.ida.plugin.item import CapaExplorerFunctionItem
+from capa.ida.plugin.model import CapaExplorerDataModel
+
+MAX_SECTION_SIZE = 750
+
+
+class CapaExplorerQtreeView(QtWidgets.QTreeView):
+    """tree view used to display hierarchical capa results
+
+    view controls UI action responses and displays data from CapaExplorerDataModel
+
+    view does not modify CapaExplorerDataModel directly - data modifications should be implemented
+    in CapaExplorerDataModel
+    """
+
+    def __init__(self, model, parent=None):
+        """initialize view"""
+        super(CapaExplorerQtreeView, self).__init__(parent)
+
+        self.setModel(model)
+
+        self.model = model
+        self.parent = parent
+
+        # control when we resize columns
+        self.should_resize_columns = True
+
+        # configure custom UI controls
+        self.setContextMenuPolicy(QtCore.Qt.CustomContextMenu)
+        self.setExpandsOnDoubleClick(False)
+        self.setSortingEnabled(True)
+        self.model.setDynamicSortFilter(False)
+
+        # configure view columns to auto-resize
+        for idx in range(CapaExplorerDataModel.COLUMN_COUNT):
+            self.header().setSectionResizeMode(idx, QtWidgets.QHeaderView.Interactive)
+
+        # disable stretch to enable horizontal scroll for last column, when needed
+        self.header().setStretchLastSection(False)
+
+        # connect slots to resize columns when expanded or collapsed
+        self.expanded.connect(self.slot_resize_columns_to_content)
+        self.collapsed.connect(self.slot_resize_columns_to_content)
+
+        # connect slots
+        self.customContextMenuRequested.connect(self.slot_custom_context_menu_requested)
+        self.doubleClicked.connect(self.slot_double_click)
+
+        self.setStyleSheet("QTreeView::item {padding-right: 15 px;padding-bottom: 2 px;}")
+
+    def reset_ui(self, should_sort=True):
+        """reset user interface changes
+
+        called when view should reset UI display e.g. expand items, resize columns
+
+        @param should_sort: True, sort results after reset, False don't sort results after reset
+        """
+        if should_sort:
+            self.sortByColumn(CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION, QtCore.Qt.AscendingOrder)
+
+        self.should_resize_columns = False
+        self.expandToDepth(0)
+        self.should_resize_columns = True
+
+        self.slot_resize_columns_to_content()
+
+    def slot_resize_columns_to_content(self):
+        """reset view columns to contents"""
+        if self.should_resize_columns:
+            self.header().resizeSections(QtWidgets.QHeaderView.ResizeToContents)
+
+            # limit size of first section
+            if self.header().sectionSize(0) > MAX_SECTION_SIZE:
+                self.header().resizeSection(0, MAX_SECTION_SIZE)
+
+    def map_index_to_source_item(self, model_index):
+        """map proxy model index to source model item
+
+        @param model_index: QModelIndex
+
+        @retval QObject
+        """
+        # assume that self.model here is either:
+        #  - CapaExplorerDataModel, or
+        #  - QSortFilterProxyModel subclass
+        #
+        # The ProxyModels may be chained,
+        #  so keep resolving the index the CapaExplorerDataModel.
+
+        model = self.model
+        while not isinstance(model, CapaExplorerDataModel):
+            if not model_index.isValid():
+                raise ValueError("invalid index")
+
+            model_index = model.mapToSource(model_index)
+            model = model.sourceModel()
+
+        if not model_index.isValid():
+            raise ValueError("invalid index")
+        return model_index.internalPointer()
+
+    def send_data_to_clipboard(self, data):
+        """copy data to the clipboard
+
+        @param data: data to be copied
+        """
+        clip = QtWidgets.QApplication.clipboard()
+        clip.clear(mode=clip.Clipboard)
+        clip.setText(data, mode=clip.Clipboard)
+
+    def new_action(self, display, data, slot):
+        """create action for context menu
+
+        @param display: text displayed to user in context menu
+        @param data: data passed to slot
+        @param slot: slot to connect
+
+        @retval QAction
+        """
+        action = QtWidgets.QAction(display, self.parent)
+        action.setData(data)
+        action.triggered.connect(lambda checked: slot(action))
+
+        return action
+
+    def load_default_context_menu_actions(self, data):
+        """yield actions specific to function custom context menu
+
+        @param data: tuple
+
+        @yield QAction
+        """
+        default_actions = (
+            ("Copy column", data, self.slot_copy_column),
+            ("Copy row", data, self.slot_copy_row),
+        )
+
+        # add default actions
+        for action in default_actions:
+            yield self.new_action(*action)
+
+    def load_function_context_menu_actions(self, data):
+        """yield actions specific to function custom context menu
+
+        @param data: tuple
+
+        @yield QAction
+        """
+        function_actions = (("Rename function", data, self.slot_rename_function),)
+
+        # add function actions
+        for action in function_actions:
+            yield self.new_action(*action)
+
+        # add default actions
+        for action in self.load_default_context_menu_actions(data):
+            yield action
+
+    def load_default_context_menu(self, pos, item, model_index):
+        """create default custom context menu
+
+        creates custom context menu containing default actions
+
+        @param pos: cursor position
+        @param item: CapaExplorerDataItem
+        @param model_index: QModelIndex
+
+        @retval QMenu
+        """
+        menu = QtWidgets.QMenu()
+
+        for action in self.load_default_context_menu_actions((pos, item, model_index)):
+            menu.addAction(action)
+
+        return menu
+
+    def load_function_item_context_menu(self, pos, item, model_index):
+        """create function custom context menu
+
+        creates custom context menu with both default actions and function actions
+
+        @param pos: cursor position
+        @param item: CapaExplorerDataItem
+        @param model_index: QModelIndex
+
+        @retval QMenu
+        """
+        menu = QtWidgets.QMenu()
+
+        for action in self.load_function_context_menu_actions((pos, item, model_index)):
+            menu.addAction(action)
+
+        return menu
+
+    def show_custom_context_menu(self, menu, pos):
+        """display custom context menu in view
+
+        @param menu: QMenu to display
+        @param pos: cursor position
+        """
+        if menu:
+            menu.exec_(self.viewport().mapToGlobal(pos))
+
+    def slot_copy_column(self, action):
+        """slot connected to custom context menu
+
+        allows user to select a column and copy the data to clipboard
+
+        @param action: QAction
+        """
+        _, item, model_index = action.data()
+        self.send_data_to_clipboard(item.data(model_index.column()))
+
+    def slot_copy_row(self, action):
+        """slot connected to custom context menu
+
+        allows user to select a row and copy the space-delimited data to clipboard
+
+        @param action: QAction
+        """
+        _, item, _ = action.data()
+        self.send_data_to_clipboard(str(item))
+
+    def slot_rename_function(self, action):
+        """slot connected to custom context menu
+
+        allows user to select a edit a function name and push changes to IDA
+
+        @param action: QAction
+        """
+        _, item, model_index = action.data()
+
+        # make item temporary edit, reset after user is finished
+        item.setIsEditable(True)
+        self.edit(model_index)
+        item.setIsEditable(False)
+
+    def slot_custom_context_menu_requested(self, pos):
+        """slot connected to custom context menu request
+
+        displays custom context menu to user containing action relevant to the item selected
+
+        @param pos: cursor position
+        """
+        model_index = self.indexAt(pos)
+        item = self.map_index_to_source_item(model_index)
+
+        column = model_index.column()
+        menu = None
+
+        if CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION == column and isinstance(item, CapaExplorerFunctionItem):
+            # user hovered function item
+            menu = self.load_function_item_context_menu(pos, item, model_index)
+        else:
+            # user hovered default item
+            menu = self.load_default_context_menu(pos, item, model_index)
+
+        # show custom context menu at view position
+        self.show_custom_context_menu(menu, pos)
+
+    def slot_double_click(self, model_index):
+        """slot connected to double-click event
+
+        if address column clicked, navigate IDA to address, else un/expand item clicked
+
+        @param model_index: QModelIndex
+        """
+        if not model_index.isValid():
+            return
+
+        item = self.map_index_to_source_item(model_index)
+        column = model_index.column()
+
+        if CapaExplorerDataModel.COLUMN_INDEX_VIRTUAL_ADDRESS == column and item.location:
+            # user double-clicked virtual address column - navigate IDA to address
+            idc.jumpto(item.location)
+
+        if CapaExplorerDataModel.COLUMN_INDEX_RULE_INFORMATION == column:
+            # user double-clicked information column - un/expand
+            self.collapse(model_index) if self.isExpanded(model_index) else self.expand(model_index)

capa/ida/plugin_helpers.py

@@ -1,93 +0,0 @@
-# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
-
-import os
-import logging
-
-import idc
-import idaapi
-from PyQt5.QtCore import Qt
-from PyQt5.QtWidgets import QTreeWidgetItem, QTreeWidgetItemIterator
-
-CAPA_EXTENSION = ".capas"
-
-
-logger = logging.getLogger("capa_ida")
-
-
-def get_input_file(freeze=True):
-    """
-    get input file path
-
-        freeze (bool): if True, get freeze file if it exists
-    """
-    # try original file in same directory as idb/i64 without idb/i64 file extension
-    input_file = idc.get_idb_path()[:-4]
-
-    if freeze:
-        # use frozen file if it exists
-        freeze_file_cand = "%s%s" % (input_file, CAPA_EXTENSION)
-        if os.path.isfile(freeze_file_cand):
-            return freeze_file_cand
-
-    if not os.path.isfile(input_file):
-        # TM naming
-        input_file = "%s.mal_" % idc.get_idb_path()[:-4]
-        if not os.path.isfile(input_file):
-            input_file = idaapi.ask_file(0, "*.*", "Please specify input file.")
-    if not input_file:
-        raise ValueError("could not find input file")
-    return input_file
-
-
-def get_orig_color_feature_vas(vas):
-    orig_colors = {}
-    for va in vas:
-        orig_colors[va] = idc.get_color(va, idc.CIC_ITEM)
-    return orig_colors
-
-
-def reset_colors(orig_colors):
-    if orig_colors:
-        for va, color in orig_colors.iteritems():
-            idc.set_color(va, idc.CIC_ITEM, orig_colors[va])
-
-
-def reset_selection(tree):
-    iterator = QTreeWidgetItemIterator(tree, QTreeWidgetItemIterator.Checked)
-    while iterator.value():
-        item = iterator.value()
-        item.setCheckState(0, Qt.Unchecked)  # column, state
-        iterator += 1
-
-
-def get_disasm_line(va):
-    return idc.generate_disasm_line(va, idc.GENDSM_FORCE_CODE)
-
-
-def get_selected_items(tree, skip_level_1=False):
-    selected = []
-    iterator = QTreeWidgetItemIterator(tree, QTreeWidgetItemIterator.Checked)
-    while iterator.value():
-        item = iterator.value()
-        if skip_level_1:
-            # hacky way to check if item is at level 1, if so, skip
-            # alternative, check if text in disasm column
-            if item.parent() and item.parent().parent() is None:
-                iterator += 1
-                continue
-        if item.text(1):
-            # logger.debug('selected %s, %s', item.text(0), item.text(1))
-            selected.append(int(item.text(1), 0x10))
-        iterator += 1
-    return selected
-
-
-def add_child_item(parent, values, feature=None):
-    child = QTreeWidgetItem(parent)
-    child.setFlags(child.flags() | Qt.ItemIsTristate | Qt.ItemIsUserCheckable)
-    for i, v in enumerate(values):
-        child.setText(i, v)
-        if feature:
-            child.setData(0, 0x100, feature)
-        child.setCheckState(0, Qt.Unchecked)
-    return child

capa/main.py

@@ -1,20 +1,25 @@
 #!/usr/bin/env python2
 """
-identify capabilities in programs.
-
 Copyright (C) 2020 FireEye, 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 os
 import sys
 import hashlib
 import logging
 import os.path
+import argparse
 import datetime
 import textwrap
 import collections
 
+import halo
 import tqdm
-import argparse
 import colorama
 
 import capa.rules
@@ -100,9 +105,14 @@ def find_capabilities(ruleset, extractor, disable_progress=None):
     all_function_matches = collections.defaultdict(list)
     all_bb_matches = collections.defaultdict(list)
 
-    meta = {"feature_counts": {"file": 0, "functions": {},}}
+    meta = {
+        "feature_counts": {
+            "file": 0,
+            "functions": {},
+        }
+    }
 
-    for f in tqdm.tqdm(extractor.get_functions(), disable=disable_progress, unit=" functions"):
+    for f in tqdm.tqdm(list(extractor.get_functions()), disable=disable_progress, desc="matching", unit=" functions"):
         function_matches, bb_matches, feature_count = find_function_capabilities(ruleset, extractor, f)
         meta["feature_counts"]["functions"][f.__int__()] = feature_count
         logger.debug("analyzed function 0x%x and extracted %d features", f.__int__(), feature_count)
@@ -206,7 +216,7 @@ def is_supported_file_type(sample):
 SHELLCODE_BASE = 0x690000
 
 
-def get_shellcode_vw(sample, arch="auto"):
+def get_shellcode_vw(sample, arch="auto", should_save=True):
     """
     Return shellcode workspace using explicit arch or via auto detect
     """
@@ -218,12 +228,17 @@ def get_shellcode_vw(sample, arch="auto"):
         # choose arch with most functions, idea by Jay G.
         vw_cands = []
         for arch in ["i386", "amd64"]:
-            vw_cands.append(viv_utils.getShellcodeWorkspace(sample_bytes, arch, base=SHELLCODE_BASE))
+            vw_cands.append(
+                viv_utils.getShellcodeWorkspace(sample_bytes, arch, base=SHELLCODE_BASE, should_save=should_save)
+            )
         if not vw_cands:
             raise ValueError("could not generate vivisect workspace")
         vw = max(vw_cands, key=lambda vw: len(vw.getFunctions()))
     else:
-        vw = viv_utils.getShellcodeWorkspace(sample_bytes, arch, base=SHELLCODE_BASE)
+        vw = viv_utils.getShellcodeWorkspace(sample_bytes, arch, base=SHELLCODE_BASE, should_save=should_save)
+
+    vw.setMeta("StorageName", "%s.viv" % sample)
+
     return vw
 
 
@@ -242,28 +257,36 @@ class UnsupportedFormatError(ValueError):
     pass
 
 
-def get_workspace(path, format):
+def get_workspace(path, format, should_save=True):
     import viv_utils
 
     logger.debug("generating vivisect workspace for: %s", path)
     if format == "auto":
         if not is_supported_file_type(path):
             raise UnsupportedFormatError()
-        vw = viv_utils.getWorkspace(path)
+        vw = viv_utils.getWorkspace(path, should_save=should_save)
     elif format == "pe":
-        vw = viv_utils.getWorkspace(path)
+        vw = viv_utils.getWorkspace(path, should_save=should_save)
     elif format == "sc32":
-        vw = get_shellcode_vw(path, arch="i386")
+        vw = get_shellcode_vw(path, arch="i386", should_save=should_save)
     elif format == "sc64":
-        vw = get_shellcode_vw(path, arch="amd64")
+        vw = get_shellcode_vw(path, arch="amd64", should_save=should_save)
     logger.debug("%s", get_meta_str(vw))
     return vw
 
 
-def get_extractor_py2(path, format):
+def get_extractor_py2(path, format, disable_progress=False):
     import capa.features.extractors.viv
 
-    vw = get_workspace(path, format)
+    with halo.Halo(text="analyzing program", spinner="simpleDots", stream=sys.stderr, enabled=not disable_progress):
+        vw = get_workspace(path, format, should_save=False)
+
+        try:
+            vw.saveWorkspace()
+        except IOError:
+            # see #168 for discussion around how to handle non-writable directories
+            logger.info("source directory is not writable, won't save intermediate workspace")
+
     return capa.features.extractors.viv.VivisectFeatureExtractor(vw, path)
 
 
@@ -271,19 +294,19 @@ class UnsupportedRuntimeError(RuntimeError):
     pass
 
 
-def get_extractor_py3(path, format):
+def get_extractor_py3(path, format, disable_progress=False):
     raise UnsupportedRuntimeError()
 
 
-def get_extractor(path, format):
+def get_extractor(path, format, disable_progress=False):
     """
     raises:
       UnsupportedFormatError:
     """
     if sys.version_info >= (3, 0):
-        return get_extractor_py3(path, format)
+        return get_extractor_py3(path, format, disable_progress=disable_progress)
     else:
-        return get_extractor_py2(path, format)
+        return get_extractor_py2(path, format, disable_progress=disable_progress)
 
 
 def is_nursery_rule_path(path):
@@ -299,7 +322,7 @@ def is_nursery_rule_path(path):
     return "nursery" in path
 
 
-def get_rules(rule_path):
+def get_rules(rule_path, disable_progress=False):
     if not os.path.exists(rule_path):
         raise IOError("rule path %s does not exist or cannot be accessed" % rule_path)
 
@@ -309,9 +332,15 @@ def get_rules(rule_path):
     elif os.path.isdir(rule_path):
         logger.debug("reading rules from directory %s", rule_path)
         for root, dirs, files in os.walk(rule_path):
+            if ".github" in root:
+                # the .github directory contains CI config in capa-rules
+                # this includes some .yml files
+                # these are not rules
+                continue
+
             for file in files:
                 if not file.endswith(".yml"):
-                    if not (file.endswith(".md") or file.endswith(".git")):
+                    if not (file.endswith(".md") or file.endswith(".git") or file.endswith(".txt")):
                         # expect to see readme.md, format.md, and maybe a .git directory
                         # other things maybe are rules, but are mis-named.
                         logger.warning("skipping non-.yml file: %s", file)
@@ -321,7 +350,8 @@ def get_rules(rule_path):
                 rule_paths.append(rule_path)
 
     rules = []
-    for rule_path in rule_paths:
+
+    for rule_path in tqdm.tqdm(list(rule_paths), disable=disable_progress, desc="loading ", unit="     rules"):
         try:
             rule = capa.rules.Rule.from_yaml_file(rule_path)
         except capa.rules.InvalidRule:
@@ -384,6 +414,7 @@ def main(argv=None):
     ]
     format_help = ", ".join(["%s: %s" % (f[0], f[1]) for f in formats])
 
+    desc = "The FLARE team's open-source tool to identify capabilities in executable files."
     epilog = textwrap.dedent(
         """
         By default, capa uses a default set of embedded rules.
@@ -396,13 +427,13 @@ def main(argv=None):
 
         examples:
           identify capabilities in a binary
-            capa suspicous.exe
+            capa suspicious.exe
 
           identify capabilities in 32-bit shellcode, see `-f` for all supported formats
             capa -f sc32 shellcode.bin
 
           report match locations
-            capa -v suspicous.exe
+            capa -v suspicious.exe
 
           report all feature match details
             capa -vv suspicious.exe
@@ -413,7 +444,7 @@ def main(argv=None):
     )
 
     parser = argparse.ArgumentParser(
-        description=__doc__, epilog=epilog, formatter_class=argparse.RawDescriptionHelpFormatter
+        description=desc, epilog=epilog, formatter_class=argparse.RawDescriptionHelpFormatter
     )
     parser.add_argument("sample", type=str, help="path to sample to analyze")
     parser.add_argument("--version", action="version", version="%(prog)s {:s}".format(capa.version.__version__))
@@ -503,7 +534,7 @@ def main(argv=None):
         logger.debug("using rules path: %s", rules_path)
 
     try:
-        rules = get_rules(rules_path)
+        rules = get_rules(rules_path, disable_progress=args.quiet)
         rules = capa.rules.RuleSet(rules)
         logger.debug("successfully loaded %s rules", len(rules))
         if args.tag:
@@ -523,7 +554,7 @@ def main(argv=None):
     else:
         format = args.format
         try:
-            extractor = get_extractor(args.sample, args.format)
+            extractor = get_extractor(args.sample, args.format, disable_progress=args.quiet)
         except UnsupportedFormatError:
             logger.error("-" * 80)
             logger.error(" Input file does not appear to be a PE file.")
@@ -547,7 +578,7 @@ def main(argv=None):
 
     meta = collect_metadata(argv, args.sample, args.rules, format, extractor)
 
-    capabilities, counts = find_capabilities(rules, extractor)
+    capabilities, counts = find_capabilities(rules, extractor, disable_progress=args.quiet)
     meta["analysis"].update(counts)
 
     if has_file_limitation(rules, capabilities):
@@ -591,6 +622,9 @@ def ida_main():
     logging.basicConfig(level=logging.INFO)
     logging.getLogger().setLevel(logging.INFO)
 
+    if not capa.ida.helpers.is_supported_ida_version():
+        return -1
+
     if not capa.ida.helpers.is_supported_file_type():
         return -1
 
@@ -613,7 +647,7 @@ def ida_main():
     rules = get_rules(rules_path)
     rules = capa.rules.RuleSet(rules)
 
-    meta = collect_metadata([], "", rules_path, format, "IdaExtractor")
+    meta = capa.ida.helpers.collect_metadata()
 
     capabilities, counts = find_capabilities(rules, capa.features.extractors.ida.IdaFeatureExtractor())
     meta["analysis"].update(counts)
@@ -621,6 +655,7 @@ def ida_main():
     if has_file_limitation(rules, capabilities, is_standalone=False):
         capa.ida.helpers.inform_user_ida_ui("capa encountered warnings during analysis")
 
+    colorama.init(strip=True)
     print(capa.render.render_default(meta, rules, capabilities))
 
 

capa/render/__init__.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 json
 
@@ -10,75 +16,59 @@ import capa.engine
 
 def convert_statement_to_result_document(statement):
     """
-        "statement": {
-            "type": "or"
-        },
+    "statement": {
+        "type": "or"
+    },
 
-        "statement": {
-            "max": 9223372036854775808,
-            "min": 2,
-            "type": "range"
-        },
+    "statement": {
+        "max": 9223372036854775808,
+        "min": 2,
+        "type": "range"
+    },
     """
-    if isinstance(statement, capa.engine.And):
-        return {
-            "type": "and",
-        }
-    elif isinstance(statement, capa.engine.Or):
-        return {
-            "type": "or",
-        }
-    elif isinstance(statement, capa.engine.Not):
-        return {
-            "type": "not",
-        }
-    elif isinstance(statement, capa.engine.Some) and statement.count == 0:
-        return {"type": "optional"}
-    elif isinstance(statement, capa.engine.Some) and statement.count > 0:
-        return {
-            "type": "some",
-            "count": statement.count,
-        }
-    elif isinstance(statement, capa.engine.Range):
-        return {
-            "type": "range",
-            "min": statement.min,
-            "max": statement.max,
-            "child": convert_feature_to_result_document(statement.child),
-        }
-    elif isinstance(statement, capa.engine.Subscope):
-        return {
-            "type": "subscope",
-            "subscope": statement.scope,
-        }
-    else:
-        raise RuntimeError("unexpected match statement type: " + str(statement))
+    statement_type = statement.name.lower()
+    result = {"type": statement_type}
+    if statement.description:
+        result["description"] = statement.description
+
+    if statement_type == "some" and statement.count == 0:
+        result["type"] = "optional"
+    elif statement_type == "some":
+        result["count"] = statement.count
+    elif statement_type == "range":
+        result["min"] = statement.min
+        result["max"] = statement.max
+        result["child"] = convert_feature_to_result_document(statement.child)
+    elif statement_type == "subscope":
+        result["subscope"] = statement.scope
+
+    return result
 
 
 def convert_feature_to_result_document(feature):
     """
-        "feature": {
-            "number": 6,
-            "type": "number"
-        },
+    "feature": {
+        "number": 6,
+        "type": "number"
+    },
 
-        "feature": {
-            "api": "ws2_32.WSASocket",
-            "type": "api"
-        },
+    "feature": {
+        "api": "ws2_32.WSASocket",
+        "type": "api"
+    },
 
-        "feature": {
-            "match": "create TCP socket",
-            "type": "match"
-        },
+    "feature": {
+        "match": "create TCP socket",
+        "type": "match"
+    },
 
-        "feature": {
-            "characteristic": [
-                "loop",
-                true
-            ],
-            "type": "characteristic"
-        },
+    "feature": {
+        "characteristic": [
+            "loop",
+            true
+        ],
+        "type": "characteristic"
+    },
     """
     result = {"type": feature.name, feature.name: feature.get_value_str()}
     if feature.description:
@@ -90,15 +80,15 @@ def convert_feature_to_result_document(feature):
 
 def convert_node_to_result_document(node):
     """
-        "node": {
-            "type": "statement",
-            "statement": { ... }
-        },
+    "node": {
+        "type": "statement",
+        "statement": { ... }
+    },
 
-        "node": {
-            "type": "feature",
-            "feature": { ... }
-        },
+    "node": {
+        "type": "feature",
+        "feature": { ... }
+    },
     """
 
     if isinstance(node, capa.engine.Statement):
@@ -162,7 +152,10 @@ def convert_match_to_result_document(rules, capabilities, result):
             scope = rule.meta["scope"]
             doc["node"] = {
                 "type": "statement",
-                "statement": {"type": "subscope", "subscope": scope,},
+                "statement": {
+                    "type": "subscope",
+                    "subscope": scope,
+                },
             }
 
         for location in doc["locations"]:
@@ -245,8 +238,8 @@ def render_verbose(meta, rules, capabilities):
 
 def render_default(meta, rules, capabilities):
     # break import loop
-    import capa.render.verbose
     import capa.render.default
+    import capa.render.verbose
 
     doc = convert_capabilities_to_result_document(meta, rules, capabilities)
     return capa.render.default.render_default(doc)
@@ -267,5 +260,7 @@ class CapaJsonObjectEncoder(json.JSONEncoder):
 
 def render_json(meta, rules, capabilities):
     return json.dumps(
-        convert_capabilities_to_result_document(meta, rules, capabilities), cls=CapaJsonObjectEncoder, sort_keys=True,
+        convert_capabilities_to_result_document(meta, rules, capabilities),
+        cls=CapaJsonObjectEncoder,
+        sort_keys=True,
     )

capa/render/default.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 collections
 
@@ -21,6 +27,8 @@ def width(s, character_count):
 def render_meta(doc, ostream):
     rows = [
         (width("md5", 22), width(doc["meta"]["sample"]["md5"], 82)),
+        ("sha1", doc["meta"]["sample"]["sha1"]),
+        ("sha256", doc["meta"]["sample"]["sha256"]),
         ("path", doc["meta"]["sample"]["path"]),
     ]
 
@@ -28,6 +36,34 @@ def render_meta(doc, ostream):
     ostream.write("\n")
 
 
+def find_subrule_matches(doc):
+    """
+    collect the rule names that have been matched as a subrule match.
+    this way we can avoid displaying entries for things that are too specific.
+    """
+    matches = set([])
+
+    def rec(node):
+        if not node["success"]:
+            # there's probably a bug here for rules that do `not: match: ...`
+            # but we don't have any examples of this yet
+            return
+
+        elif node["node"]["type"] == "statement":
+            for child in node["children"]:
+                rec(child)
+
+        elif node["node"]["type"] == "feature":
+            if node["node"]["feature"]["type"] == "match":
+                matches.add(node["node"]["feature"]["match"])
+
+    for rule in rutils.capability_rules(doc):
+        for node in rule["matches"].values():
+            rec(node)
+
+    return matches
+
+
 def render_capabilities(doc, ostream):
     """
     example::
@@ -40,8 +76,16 @@ def render_capabilities(doc, ostream):
         | ...                                                   | ...                                             |
         +-------------------------------------------------------+-------------------------------------------------+
     """
+    subrule_matches = find_subrule_matches(doc)
+
     rows = []
     for rule in rutils.capability_rules(doc):
+        if rule["meta"]["name"] in subrule_matches:
+            # rules that are also matched by other rules should not get rendered by default.
+            # this cuts down on the amount of output while giving approx the same detail.
+            # see #224
+            continue
+
         count = len(rule["matches"])
         if count == 1:
             capability = rutils.bold(rule["meta"]["name"])
@@ -101,7 +145,12 @@ def render_attack(doc, ostream):
                 inner_rows.append("%s::%s %s" % (rutils.bold(technique), subtechnique, id))
             else:
                 raise RuntimeError("unexpected ATT&CK spec format")
-        rows.append((rutils.bold(tactic.upper()), "\n".join(inner_rows),))
+        rows.append(
+            (
+                rutils.bold(tactic.upper()),
+                "\n".join(inner_rows),
+            )
+        )
 
     if rows:
         ostream.write(

capa/render/utils.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 six
 import termcolor
@@ -16,7 +22,10 @@ def bold2(s):
 
 def hex(n):
     """render the given number using upper case hex, like: 0x123ABC"""
-    return "0x%X" % n
+    if n < 0:
+        return "-0x%X" % (-n)
+    else:
+        return "0x%X" % n
 
 
 def capability_rules(doc):

capa/render/verbose.py

@@ -15,6 +15,12 @@ example::
                  0x10003797
 
 Copyright (C) 2020 FireEye, 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 tabulate
 

capa/render/vverbose.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 collections
 
@@ -30,15 +36,17 @@ def render_locations(ostream, match):
 
 def render_statement(ostream, match, statement, indent=0):
     ostream.write("  " * indent)
-    if statement["type"] in ("and", "or", "optional"):
+    if statement["type"] in ("and", "or", "optional", "not", "subscope"):
         ostream.write(statement["type"])
-        ostream.writeln(":")
-    elif statement["type"] == "not":
-        # this statement is handled specially in `render_match` using the MODE_SUCCESS/MODE_FAILURE flags.
-        ostream.writeln("not:")
+        ostream.write(":")
+        if statement.get("description"):
+            ostream.write(" = %s" % statement["description"])
+        ostream.writeln("")
     elif statement["type"] == "some":
-        ostream.write(statement["count"] + " or more")
-        ostream.writeln(":")
+        ostream.write("%d or more:" % (statement["count"]))
+        if statement.get("description"):
+            ostream.write(" = %s" % statement["description"])
+        ostream.writeln("")
     elif statement["type"] == "range":
         # `range` is a weird node, its almost a hybrid of statement+feature.
         # it is a specific feature repeated multiple times.
@@ -65,11 +73,10 @@ def render_statement(ostream, match, statement, indent=0):
         else:
             ostream.write("between %d and %d" % (statement["min"], statement["max"]))
 
+        if statement.get("description"):
+            ostream.write(" = %s" % statement["description"])
         render_locations(ostream, match)
-        ostream.write("\n")
-    elif statement["type"] == "subscope":
-        ostream.write(statement["subscope"])
-        ostream.writeln(":")
+        ostream.writeln("")
     else:
         raise RuntimeError("unexpected match statement type: " + str(statement))
 
@@ -198,7 +205,7 @@ def render_rules(ostream, doc):
                 # because we do the file-scope evaluation a single time.
                 # but i'm not 100% sure if this is/will always be true.
                 # so, lets be explicit about our assumptions and raise an exception if they fail.
-                raise RuntimeError("unexpected file scope match count: " + len(matches))
+                raise RuntimeError("unexpected file scope match count: %d" % (len(matches)))
             render_match(ostream, matches[0], indent=0)
         else:
             for location, match in sorted(doc["rules"][rule["meta"]["name"]]["matches"].items()):

capa/rules.py

@@ -1,9 +1,16 @@
 # Copyright (C) 2020 FireEye, 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 uuid
 import codecs
 import logging
 import binascii
+import functools
 
 import six
 import ruamel.yaml
@@ -62,7 +69,6 @@ SUPPORTED_FEATURES = {
     FUNCTION_SCOPE: {
         # plus basic block scope features, see below
         capa.features.basicblock.BasicBlock,
-        capa.features.Characteristic("switch"),
         capa.features.Characteristic("calls from"),
         capa.features.Characteristic("calls to"),
         capa.features.Characteristic("loop"),
@@ -189,8 +195,20 @@ def parse_feature(key):
         return capa.features.Bytes
     elif key == "number":
         return capa.features.insn.Number
+    elif key.startswith("number/"):
+        arch = key.partition("/")[2]
+        # the other handlers here return constructors for features,
+        # and we want to as well,
+        # however, we need to preconfigure one of the arguments (`arch`).
+        # so, instead we return a partially-applied function that
+        #  provides `arch` to the feature constructor.
+        # it forwards any other arguments provided to the closure along to the constructor.
+        return functools.partial(capa.features.insn.Number, arch=arch)
     elif key == "offset":
         return capa.features.insn.Offset
+    elif key.startswith("offset/"):
+        arch = key.partition("/")[2]
+        return functools.partial(capa.features.insn.Offset, arch=arch)
     elif key == "mnemonic":
         return capa.features.insn.Mnemonic
     elif key == "basic blocks":
@@ -244,7 +262,7 @@ def parse_description(s, value_type, description=None):
                 raise InvalidRule(
                     "unexpected bytes value: byte sequences must be no larger than %s bytes" % MAX_BYTES_FEATURE_SIZE
                 )
-        elif value_type in {"number", "offset"}:
+        elif value_type in ("number", "offset") or value_type.startswith(("number/", "offset/")):
             try:
                 value = parse_int(value)
             except ValueError:
@@ -259,21 +277,21 @@ def build_statements(d, scope):
 
     key = list(d.keys())[0]
     if key == "and":
-        return And(*[build_statements(dd, scope) for dd in d[key]])
+        return And([build_statements(dd, scope) for dd in d[key]], description=d.get("description"))
     elif key == "or":
-        return Or(*[build_statements(dd, scope) for dd in d[key]])
+        return Or([build_statements(dd, scope) for dd in d[key]], description=d.get("description"))
     elif key == "not":
         if len(d[key]) != 1:
             raise InvalidRule("not statement must have exactly one child statement")
-        return Not(*[build_statements(dd, scope) for dd in d[key]])
+        return Not(build_statements(d[key][0], scope), description=d.get("description"))
     elif key.endswith(" or more"):
         count = int(key[: -len("or more")])
-        return Some(count, *[build_statements(dd, scope) for dd in d[key]])
+        return Some(count, [build_statements(dd, scope) for dd in d[key]], description=d.get("description"))
     elif key == "optional":
         # `optional` is an alias for `0 or more`
         # which is useful for documenting behaviors,
         # like with `write file`, we might say that `WriteFile` is optionally found alongside `CreateFileA`.
-        return Some(0, *[build_statements(dd, scope) for dd in d[key]])
+        return Some(0, [build_statements(dd, scope) for dd in d[key]], description=d.get("description"))
 
     elif key == "function":
         if scope != FILE_SCOPE:
@@ -282,7 +300,7 @@ def build_statements(d, scope):
         if len(d[key]) != 1:
             raise InvalidRule("subscope must have exactly one child statement")
 
-        return Subscope(FUNCTION_SCOPE, *[build_statements(dd, FUNCTION_SCOPE) for dd in d[key]])
+        return Subscope(FUNCTION_SCOPE, build_statements(d[key][0], FUNCTION_SCOPE))
 
     elif key == "basic block":
         if scope != FUNCTION_SCOPE:
@@ -291,7 +309,7 @@ def build_statements(d, scope):
         if len(d[key]) != 1:
             raise InvalidRule("subscope must have exactly one child statement")
 
-        return Subscope(BASIC_BLOCK_SCOPE, *[build_statements(dd, BASIC_BLOCK_SCOPE) for dd in d[key]])
+        return Subscope(BASIC_BLOCK_SCOPE, build_statements(d[key][0], BASIC_BLOCK_SCOPE))
 
     elif key.startswith("count(") and key.endswith(")"):
         # e.g.:
@@ -319,7 +337,7 @@ def build_statements(d, scope):
             #     count(number(0x100 = description))
             if term != "string":
                 value, description = parse_description(arg, term)
-                feature = Feature(value, description)
+                feature = Feature(value, description=description)
             else:
                 # arg is string (which doesn't support inline descriptions), like:
                 #
@@ -332,18 +350,18 @@ def build_statements(d, scope):
 
         count = d[key]
         if isinstance(count, int):
-            return Range(feature, min=count, max=count)
+            return Range(feature, min=count, max=count, description=d.get("description"))
         elif count.endswith(" or more"):
             min = parse_int(count[: -len(" or more")])
             max = None
-            return Range(feature, min=min, max=max)
+            return Range(feature, min=min, max=max, description=d.get("description"))
         elif count.endswith(" or fewer"):
             min = None
             max = parse_int(count[: -len(" or fewer")])
-            return Range(feature, min=min, max=max)
+            return Range(feature, min=min, max=max, description=d.get("description"))
         elif count.startswith("("):
             min, max = parse_range(count)
-            return Range(feature, min=min, max=max)
+            return Range(feature, min=min, max=max, description=d.get("description"))
         else:
             raise InvalidRule("unexpected range: %s" % (count))
     elif key == "string" and not isinstance(d[key], six.string_types):
@@ -352,7 +370,7 @@ def build_statements(d, scope):
         Feature = parse_feature(key)
         value, description = parse_description(d[key], key, d.get("description"))
         try:
-            feature = Feature(value, description)
+            feature = Feature(value, description=description)
         except ValueError as e:
             raise InvalidRule(str(e))
         ensure_feature_valid_for_scope(scope, feature)
@@ -606,7 +624,25 @@ class Rule(object):
                 continue
             meta[key] = value
 
-        return ostream.getvalue().decode("utf-8").rstrip("\n") + "\n"
+        doc = ostream.getvalue().decode("utf-8").rstrip("\n") + "\n"
+        # when we have something like:
+        #
+        #     and:
+        #       - string: foo
+        #         description: bar
+        #
+        # we want the `description` horizontally aligned with the start of the `string` (like above).
+        # however, ruamel will give us (which I don't think is even valid yaml):
+        #
+        #     and:
+        #       - string: foo
+        #      description: bar
+        #
+        # tweaking `ruamel.indent()` doesn't quite give us the control we want.
+        # so, add the two extra spaces that we've determined we need through experimentation.
+        # see #263
+        doc = doc.replace("  description:", "    description:")
+        return doc
 
 
 def get_rules_with_scope(rules, scope):

capa/version.py

@@ -1,2 +1 @@
-__version__ = "1.0.0"
-__commit__ = "00000000"
+__version__ = "1.3.0"

doc/installation.md

@@ -8,58 +8,60 @@ We use PyInstaller to create these packages.
 
 The capa [README](../README.md#download) also links to nightly builds of standalone binaries from the latest development branch.
 
+### Linux Standalone installation
+
+The Linux Standalone binary has been built using GLIB 2.26.
+Consequently it works when using GLIB >= 2.26.
+This requirement is satisfied by default in most newer distribution such as Ubuntu >= 18, Debian >= 10, openSUSE >= 15.1 and CentOS >= 8.
+But the binary may not work in older distributions.
+
+### MacOS Standalone installation
+
+By default, on MacOS Catalina or greater, Gatekeeper will block execution of the standalone binary. To resolve this, simply try to execute it once on the command-line and then go to `System Preferences` / `Security & Privacy` / `General` and approve the application:
+
+![approve dialog](img/approve.png)
+
 ## Method 2: Using capa as a Python library
 To install capa as a Python library, you'll need to install a few dependencies, and then use `pip` to fetch the capa module.
-Note: this technique doesn't pull the default rule set, so you should check it out separately from [capa-rules](https://github.com/fireeye/capa-rules/) and pass the directory to the entrypoint using `-r`.
 
-### 1. Install requirements
-First, install the requirements.
-`$ pip install https://github.com/williballenthin/vivisect/zipball/master`
+#### *Note*:
+This method is appropriate for integrating capa in an existing project. It is not the right choice for local tool usage, such as within IDA Pro - see Method 3, instead.
+That's because this technique doesn't pull the default rule set, so you should check it out separately from [capa-rules](https://github.com/fireeye/capa-rules/) and pass the directory to the entrypoint using `-r`.
 
-### 2. Install capa module
+### 1. Install capa module
 Second, use `pip` to install the capa module to your local Python environment. This fetches the library code to your computer but does not keep editable source files around for you to hack on. If you'd like to edit the source files, see below.
 `$ pip install https://github.com/fireeye/capa/archive/master.zip`
 
-### 3. Use capa
+### 2. Use capa
 You can now import the `capa` module from a Python script or use the IDA Pro plugins from the `capa/ida` directory. For more information please see the [usage](usage.md) documentation.
 
 ## Method 3: Inspecting the capa source code
 If you'd like to review and modify the capa source code, you'll need to check it out from GitHub and install it locally. By following these instructions, you'll maintain a local directory of source code that you can modify and run easily. 
 
-### 1. Install requirements
-First, install the requirements.
-`$ pip install https://github.com/williballenthin/vivisect/zipball/master`
+### 1. Check out source code
+Next, clone the capa git repository.
+We use submodules to separate [code](https://github.com/fireeye/capa), [rules](https://github.com/fireeye/capa-rules), and [test data](https://github.com/fireeye/capa-testfiles).
+To clone everything use the `--recurse-submodules` option:
+- `$ git clone --recurse-submodules https://github.com/fireeye/capa.git /local/path/to/src` (HTTPS)
+- `$ git clone --recurse-submodules git@github.com:fireeye/capa.git /local/path/to/src` (SSH)
 
-### 2. Check out source code
-Next, clone the capa git repository. We use submodules to separate code, rules, and test data. See below to get all data at once. To only get the source code and our provided rules (common), follow these steps:
+To only get the source code and our provided rules (common), follow these steps:
 - clone repository
   - `$ git clone https://github.com/fireeye/capa.git /local/path/to/src` (HTTPS)
   - `$ git clone git@github.com:fireeye/capa.git /local/path/to/src` (SSH)
 - `$ cd /local/path/to/src`
-- `$ git submodule init`
-- `$ git submodule update rules`
+- `$ git submodule update --init rules`
 
-#### capa-testfiles
-The [capa-testfiles](https://github.com/fireeye/capa-testfiles) repository (`/local/path/to/src/tests/data`) contains a large collection of malware and benign test files. *In most cases you will not need to check it out on your local system.*
-
-To update the testfiles you can use the following command:
-- `$ git submodule update tests/data`
-
-To get all data at once use the `--recurse-submodules` option:
-
-- `$ git clone --recurse-submodules https://github.com/fireeye/capa.git /local/path/to/src` (HTTPS)
-- `$ git clone --recurse-submodules git@github.com:fireeye/capa.git /local/path/to/src` (SSH)
-
-### 3. Install the local source code
-Finally, use `pip` to install the source code in "editable" mode. This means that Python will load the capa module from the local directory rather than copying it to `site-packages` or `dist-packages`. This is good because it is easy to modify files and see the effects reflected immediately. But, be careful not to remove this directory unless uninstalling capa.
+### 2. Install the local source code
+Use `pip` to install the source code in "editable" mode. This means that Python will load the capa module from the local directory rather than copying it to `site-packages` or `dist-packages`. This is good because it is easy to modify files and see the effects reflected immediately. But, be careful not to remove this directory unless uninstalling capa.
 
 `$ pip install -e /local/path/to/src`
 
-You'll find that the `capa.exe` (Windows) or `capa` (Linux) executables in your path now invoke the capa binary from this directory.
+You'll find that the `capa.exe` (Windows) or `capa` (Linux/MacOS) executables in your path now invoke the capa binary from this directory.
 
 We use the following tools to ensure consistent code style and formatting:
   - [black](https://github.com/psf/black) code formatter, with `-l 120`
-  - [isort](https://pypi.org/project/isort/) code formatter, with `--length-sort --line-width 120`
+  - [isort 5](https://pypi.org/project/isort/) code formatter, with `--profile black --length-sort --line-width 120`
   - [dos2unix](https://linux.die.net/man/1/dos2unix) for UNIX-style LF newlines
   - [capafmt](https://github.com/fireeye/capa/blob/master/scripts/capafmt.py) rule formatter
 
@@ -69,10 +71,17 @@ To install these development dependencies, run:
 
 Note that some development dependencies (including the black code formatter) require Python 3.
 
-### 4. Setup hooks [optional]
+To check the code style, formatting and run the tests you can run the script `scripts/ci.sh`.
+You can run it with the argument `no_tests` to skip the tests and only run the code style and formatting: `scripts/ci.sh no_tests`
+
+### 3. Setup hooks [optional]
 
 If you plan to contribute to capa, you may want to setup the hooks.
 Run `scripts/setup-hooks.sh` to set the following hooks up:
-- The `post-commit` hook runs checks after every `git commit`, letting you know if there are code style or rule linter offenses you need to fix.
-- The `pre-push` hook runs various style and lint checks as well as the tests. If they do not succeed `git push` is blocked.
+- The `pre-commit` hook runs checks before every `git commit`.
+  It runs `scripts/ci.sh no_tests` aborting the commit if there are code style or rule linter offenses you need to fix.
+  You can skip this check by using the `--no-verify` git option.
+- The `pre-push` hook runs checks before every `git push`.
+  It runs `scripts/ci.sh` aborting the push if there are code style or rule linter offenses or if the tests fail.
   This way you can ensure everything is alright before sending a pull request.
+

doc/usage.md

@@ -12,29 +12,3 @@ See `capa -h` for all supported arguments and usage examples.
 Use the `-t` option to run rules with the given metadata value (see the rule fields `rule.meta.*`).
 For example, `capa -t william.ballenthin@mandiant.com` runs rules that reference Willi's email address (probably as the author), or
 `capa -t communication` runs rules with the namespace `communication`.
-
-### IDA Pro integrations
-You can run capa from within IDA Pro. Run `capa/main.py` via `File - Script file...` (or ALT + F7). 
-When running in IDA, capa uses IDA's disassembly and file analysis as its backend. 
-These results may vary from the standalone version that uses vivisect.
-IDA's analysis is generally a bit faster and more thorough than vivisect's, so you might prefer this mode.
-
-When run under IDA, capa supports both Python 2 and Python 3 interpreters.
-If you encounter issues with your specific setup, please open a new [Issue](https://github.com/fireeye/capa/issues).
-
-Additionally, capa comes with an IDA Pro plugin located in the `capa/ida` directory: the explorer.
-
-#### capa explorer
-The capa explorer allows you to interactively display and browse capabilities capa identified in a binary.
-As you select rules or logic, capa will highlight the addresses that support its analysis conclusions.
-We like to use capa to help find the most interesting parts of a program, such as where the C2 mechanism might be.
-
-![capa explorer](img/capa_explorer.png)
-
-To install the plugin, you'll need to be running IDA Pro 7.4 or 7.5 with either Python 2 or Python 3.
-Next make sure pip commands are run using the Python install that is configured for your IDA install:
-
-  1. Only if running Python 2.7, run command `$ pip install https://github.com/williballenthin/vivisect/zipball/master`
-  2. Run `$ pip install .` from capa root directory
-  3. Open IDA and navigate to `File > Script file…` or `Alt+F7`
-  4. Navigate to `<capa_install_dir>\capa\ida\` and choose `ida_capa_explorer.py`

scripts/capafmt.py

@@ -7,10 +7,15 @@ Usage:
    $ python capafmt.py -i foo.yml
 
 Copyright (C) 2020 FireEye, 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 sys
 import logging
-
 import argparse
 
 import capa.rules

scripts/ci.sh

@@ -0,0 +1,86 @@
+#!/usr/bin/env bash
+
+# Copyright (C) 2020 FireEye, 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.
+
+# Use a console with emojis support for a better experience
+
+# Stash uncommited changes
+MSG="pre-push-$(date +%s)";
+git stash push -kum "$MSG" &>/dev/null ;
+STASH_LIST=$(git stash list);
+if [[ "$STASH_LIST" == *"$MSG"* ]]; then
+  echo "Uncommited changes stashed with message '$MSG', if you abort before they are restored run \`git stash pop\`";
+fi
+
+restore_stashed() {
+  if [[ "$STASH_LIST" == *"$MSG"* ]]; then
+    git stash pop --index &>/dev/null ;
+    echo "Stashed changes '$MSG' restored";
+  fi
+}
+
+python_3() {
+  case "$(uname -s)" in
+   CYGWIN*|MINGW32*|MSYS*|MINGW*)
+     py -3 -m $1 > $2 2>&1;;
+   *)
+     python3 -m $1 > $2 2>&1;;
+  esac
+}
+
+# Run isort and print state
+python_3 'isort --profile black --length-sort --line-width 120 -c .' 'isort-output.log';
+if [ $? == 0 ]; then
+  echo 'isort succeeded!! 💖';
+else
+  echo 'isort FAILED! 😭';
+  echo 'Check isort-output.log for details';
+  restore_stashed;
+  exit 1;
+fi
+
+# Run black and print state
+python_3 'black -l 120 --check .' 'black-output.log';
+if [ $? == 0 ]; then
+  echo 'black succeeded!! 💝';
+else
+  echo 'black FAILED! 😭';
+  echo 'Check black-output.log for details';
+  restore_stashed;
+  exit 2;
+fi
+
+# Run rule linter and print state
+python ./scripts/lint.py ./rules/ > rule-linter-output.log 2>&1;
+if [ $? == 0 ]; then
+  echo 'Rule linter succeeded!! 💘';
+else
+  echo 'Rule linter FAILED! 😭';
+  echo 'Check rule-linter-output.log for details';
+  restore_stashed;
+  exit 3;
+fi
+
+# Run tests except if first argument is no_tests
+if [ "$1" != 'no_tests' ]; then
+  echo 'Running tests, please wait ⌛';
+  pytest tests/ --maxfail=1;
+  if [ $? == 0 ]; then
+    echo 'Tests succeed!! 🎉';
+  else
+    echo 'Tests FAILED! 😓';
+    echo 'Run `pytest -v --cov=capa test/` if you need more details';
+    restore_stashed;
+    exit 4;
+  fi
+fi
+
+restore_stashed;
+echo 'SUCCEEDED 🎉🎉';
+

scripts/hooks/post-commit

@@ -1,37 +0,0 @@
-#!/usr/bin/env bash
-# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
-# doesn't matter if this gets repeated later on in a hooks file
-
-# Use a console with emojis support for a better experience
-
-# Stash uncommited changes
-MSG="post-commit-$(date +%s)";
-git stash push -kqum "$MSG";
-STASH_LIST=$(git stash list);
-if [[ "$STASH_LIST" == *"$MSG"* ]]; then
-  echo "Uncommited changes stashed with message '$MSG', if you abort before they are restored run \`git stash pop\`";
-fi
-
-# Run style checker and print state (it doesn't block the commit)
-pycodestyle --config=./.github/tox.ini ./capa/ > style-checker-output.log 2>&1;
-if [ $? == 0 ]; then
-  echo 'Style checker succeeds!! 💘';
-else
-  echo 'Style checker failed 😭';
-  echo 'Check style-checker-output.log for details';
-fi
-
-# Run rule linter and print state (it doesn't block the commit)
-python ./scripts/lint.py ./rules/ > rule-linter-output.log 2>&1;
-if [ $? == 0 ]; then
-  echo 'Rule linter succeeds!! 💖';
-else
-  echo 'Rule linter failed 😭';
-  echo 'Check rule-linter-output.log for details';
-fi
-
-# Restore stashed changes
-if [[ "$STASH_LIST" == *"$MSG"* ]]; then
-  git stash pop -q --index;
-  echo "Stashed changes '$MSG' restored";
-fi

scripts/hooks/pre-push

@@ -1,58 +0,0 @@
-#!/usr/bin/env bash
-# Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
-# doesn't matter if this gets repeated later on in a hooks file
-
-# Use a console with emojis support for a better experience
-
-# Stash uncommited changes
-MSG="pre-push-$(date +%s)";
-git stash push -kqum "$MSG";
-STASH_LIST=$(git stash list);
-if [[ "$STASH_LIST" == *"$MSG"* ]]; then
-  echo "Uncommited changes stashed with message '$MSG', if you abort before they are restored run \`git stash pop\`";
-fi
-
-restore_stashed() {
-  if [[ "$STASH_LIST" == *"$MSG"* ]]; then
-    git stash pop -q --index;
-    echo "Stashed changes '$MSG' restored";
-  fi
-}
-
-# Run style checker and print state
-pycodestyle --config=./.github/tox.ini ./capa/ > style-checker-output.log 2>&1;
-if [ $? == 0 ]; then
-  echo 'Style checker succeeds!! 💘';
-else
-  echo 'Style checker failed 😭 PUSH ABORTED';
-  echo 'Check style-checker-output.log for details';
-  restore_stashed;
-  exit 1;
-fi
-
-# Run rule linter and print state
-python ./scripts/lint.py ./rules/ > rule-linter-output.log 2>&1;
-if [ $? == 0 ]; then
-  echo 'Rule linter succeeds!! 💖';
-else
-  echo 'Rule linter failed 😭 PUSH ABORTED';
-  echo 'Check rule-linter-output.log for details';
-  restore_stashed;
-  exit 2;
-fi
-
-# Run tests
-echo 'Running tests, please wait ⌛';
-pytest tests/ --maxfail=1;
-if [ $? == 0 ]; then
-  echo 'Tests succeed!! 🎉';
-else
-  echo 'Tests failed 😓 PUSH ABORTED';
-  echo 'Run `pytest -v --cov=capa test/` if you need more details';
-  restore_stashed;
-  exit 3;
-fi
-
-echo 'PUSH SUCCEEDED 🎉🎉';
-
-restore_stashed;

scripts/import-to-bn.py

@@ -0,0 +1,111 @@
+"""
+Binary Ninja plugin that imports a capa report,
+produced via `capa --json /path/to/sample`,
+into the current database.
+
+It will mark up functions with their capa matches, like:
+
+    ; capa: print debug messages (host-interaction/log/debug/write-event)
+    ; capa: delete service (host-interaction/service/delete)
+    ; Attributes: bp-based frame
+
+    public UninstallService
+    UninstallService proc near
+    ...
+
+To use, invoke from the Binary Ninja Tools menu, or from the command-palette.
+
+Adapted for Binary Ninja by @psifertex
+
+This script will verify that the report matches the workspace.
+Check the log window for any errors, and/or the summary of changes.
+
+Derived from: https://github.com/fireeye/capa/blob/master/scripts/import-to-ida.py
+"""
+import os
+import json
+
+from binaryninja import *
+
+
+def append_func_cmt(bv, va, cmt):
+    """
+    add the given comment to the given function,
+    if it doesn't already exist.
+    """
+    func = bv.get_function_at(va)
+    if not func:
+        raise ValueError("not a function")
+
+    if cmt in func.comment:
+        return
+
+    func.comment = func.comment + "\n" + cmt
+
+
+def load_analysis(bv):
+    shortname = os.path.splitext(os.path.basename(bv.file.filename))[0]
+    dirname = os.path.dirname(bv.file.filename)
+    log_info(f"dirname: {dirname}\nshortname: {shortname}\n")
+    if os.access(os.path.join(dirname, shortname + ".js"), os.R_OK):
+        path = os.path.join(dirname, shortname + ".js")
+    elif os.access(os.path.join(dirname, shortname + ".json"), os.R_OK):
+        path = os.path.join(dirname, shortname + ".json")
+    else:
+        path = interaction.get_open_filename_input("capa report:", "JSON (*.js *.json);;All Files (*)")
+    if not path or not os.access(path, os.R_OK):
+        log_error("Invalid filename.")
+        return 0
+    log_info("Using capa file %s" % path)
+
+    with open(path, "rb") as f:
+        doc = json.loads(f.read().decode("utf-8"))
+
+    if "meta" not in doc or "rules" not in doc:
+        log_error("doesn't appear to be a capa report")
+        return -1
+
+    a = doc["meta"]["sample"]["md5"].lower()
+    md5 = Transform["MD5"]
+    rawhex = Transform["RawHex"]
+    b = rawhex.encode(md5.encode(bv.parent_view.read(bv.parent_view.start, bv.parent_view.end))).decode("utf-8")
+    if not a == b:
+        log_error("sample mismatch")
+        return -2
+
+    rows = []
+    for rule in doc["rules"].values():
+        if rule["meta"].get("lib"):
+            continue
+        if rule["meta"].get("capa/subscope"):
+            continue
+        if rule["meta"]["scope"] != "function":
+            continue
+
+        name = rule["meta"]["name"]
+        ns = rule["meta"].get("namespace", "")
+        for va in rule["matches"].keys():
+            va = int(va)
+            rows.append((ns, name, va))
+
+    # order by (namespace, name) so that like things show up together
+    rows = sorted(rows)
+    for ns, name, va in rows:
+        if ns:
+            cmt = "%s (%s)" % (name, ns)
+        else:
+            cmt = "%s" % (name,)
+
+        log_info("0x%x: %s" % (va, cmt))
+        try:
+            # message will look something like:
+            #
+            #     capa: delete service (host-interaction/service/delete)
+            append_func_cmt(bv, va, "capa: " + cmt)
+        except ValueError:
+            continue
+
+    log_info("ok")
+
+
+PluginCommand.register("Load capa file", "Loads an analysis file from capa", load_analysis)

scripts/import-to-ida.py

@@ -1,111 +1,118 @@
-"""
-IDA Pro script that imports a capa report,
-produced via `capa --json /path/to/sample`,
-into the current database.
-
-It will mark up functions with their capa matches, like:
-
-    ; capa: print debug messages (host-interaction/log/debug/write-event)
-    ; capa: delete service (host-interaction/service/delete)
-    ; Attributes: bp-based frame
-
-    public UninstallService
-    UninstallService proc near
-    ...
-
-To use, invoke from the IDA Pro scripting dialog,
-such as via Alt-F9,
-and then select the existing capa report from the file system.
-
-This script will verify that the report matches the workspace.
-Check the output window for any errors, and/or the summary of changes.
-
-Copyright (C) 2020 FireEye, Inc. All Rights Reserved.
-"""
-import json
-import logging
-
-import idc
-import idautils
-import ida_idaapi
-import ida_kernwin
-
-logger = logging.getLogger("capa")
-
-
-def append_func_cmt(va, cmt, repeatable=False):
-    """
-    add the given comment to the given function, 
-    if it doesn't already exist.
-    """
-    func = ida_funcs.get_func(va)
-    if not func:
-        raise ValueError("not a function")
-
-    existing = ida_funcs.get_func_cmt(func, repeatable) or ""
-    if cmt in existing:
-        return
-
-    new = existing + "\n" + cmt
-    ida_funcs.set_func_cmt(func, new, repeatable)
-
-
-def main():
-    path = ida_kernwin.ask_file(False, "*", "capa report")
-    if not path:
-        return 0
-
-    with open(path, "rb") as f:
-        doc = json.loads(f.read().decode("utf-8"))
-
-    if "meta" not in doc or "rules" not in doc:
-        logger.error("doesn't appear to be a capa report")
-        return -1
-
-    # in IDA 7.4, the MD5 hash may be truncated, for example:
-    # wanted: 84882c9d43e23d63b82004fae74ebb61
-    # found: b'84882C9D43E23D63B82004FAE74EBB6\x00'
-    #
-    # see: https://github.com/idapython/bin/issues/11
-    a = doc["meta"]["sample"]["md5"].lower()
-    b = idautils.GetInputFileMD5().decode("ascii").lower().rstrip("\x00")
-    if not a.startswith(b):
-        logger.error("sample mismatch")
-        return -2
-
-    rows = []
-    for rule in doc["rules"].values():
-        if rule["meta"].get("lib"):
-            continue
-        if rule["meta"].get("capa/subscope"):
-            continue
-        if rule["meta"]["scope"] != "function":
-            continue
-
-        name = rule["meta"]["name"]
-        ns = rule["meta"].get("namespace", "")
-        for va in rule["matches"].keys():
-            va = int(va)
-            rows.append((ns, name, va))
-
-    # order by (namespace, name) so that like things show up together
-    rows = sorted(rows)
-    for ns, name, va in rows:
-        if ns:
-            cmt = "%s (%s)" % (name, ns)
-        else:
-            cmt = "%s" % (name,)
-
-        logger.info("0x%x: %s", va, cmt)
-        try:
-            # message will look something like:
-            #
-            #     capa: delete service (host-interaction/service/delete)
-            append_func_cmt(va, "capa: " + cmt, repeatable=False)
-        except ValueError:
-            continue
-
-    logger.info("ok")
-
-
-main()
+"""
+IDA Pro script that imports a capa report,
+produced via `capa --json /path/to/sample`,
+into the current database.
+
+It will mark up functions with their capa matches, like:
+
+    ; capa: print debug messages (host-interaction/log/debug/write-event)
+    ; capa: delete service (host-interaction/service/delete)
+    ; Attributes: bp-based frame
+
+    public UninstallService
+    UninstallService proc near
+    ...
+
+To use, invoke from the IDA Pro scripting dialog,
+such as via Alt-F9,
+and then select the existing capa report from the file system.
+
+This script will verify that the report matches the workspace.
+Check the output window for any errors, and/or the summary of changes.
+
+Copyright (C) 2020 FireEye, 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 json
+import logging
+
+import idc
+import idautils
+import ida_funcs
+import ida_idaapi
+import ida_kernwin
+
+logger = logging.getLogger("capa")
+
+
+def append_func_cmt(va, cmt, repeatable=False):
+    """
+    add the given comment to the given function,
+    if it doesn't already exist.
+    """
+    func = ida_funcs.get_func(va)
+    if not func:
+        raise ValueError("not a function")
+
+    existing = ida_funcs.get_func_cmt(func, repeatable) or ""
+    if cmt in existing:
+        return
+
+    new = existing + "\n" + cmt
+    ida_funcs.set_func_cmt(func, new, repeatable)
+
+
+def main():
+    path = ida_kernwin.ask_file(False, "*", "capa report")
+    if not path:
+        return 0
+
+    with open(path, "rb") as f:
+        doc = json.loads(f.read().decode("utf-8"))
+
+    if "meta" not in doc or "rules" not in doc:
+        logger.error("doesn't appear to be a capa report")
+        return -1
+
+    # in IDA 7.4, the MD5 hash may be truncated, for example:
+    # wanted: 84882c9d43e23d63b82004fae74ebb61
+    # found: b'84882C9D43E23D63B82004FAE74EBB6\x00'
+    #
+    # see: https://github.com/idapython/bin/issues/11
+    a = doc["meta"]["sample"]["md5"].lower()
+    b = idautils.GetInputFileMD5().decode("ascii").lower().rstrip("\x00")
+    if not a.startswith(b):
+        logger.error("sample mismatch")
+        return -2
+
+    rows = []
+    for rule in doc["rules"].values():
+        if rule["meta"].get("lib"):
+            continue
+        if rule["meta"].get("capa/subscope"):
+            continue
+        if rule["meta"]["scope"] != "function":
+            continue
+
+        name = rule["meta"]["name"]
+        ns = rule["meta"].get("namespace", "")
+        for va in rule["matches"].keys():
+            va = int(va)
+            rows.append((ns, name, va))
+
+    # order by (namespace, name) so that like things show up together
+    rows = sorted(rows)
+    for ns, name, va in rows:
+        if ns:
+            cmt = "%s (%s)" % (name, ns)
+        else:
+            cmt = "%s" % (name,)
+
+        logger.info("0x%x: %s", va, cmt)
+        try:
+            # message will look something like:
+            #
+            #     capa: delete service (host-interaction/service/delete)
+            append_func_cmt(va, "capa: " + cmt, repeatable=False)
+        except ValueError:
+            continue
+
+    logger.info("ok")
+
+
+main()

scripts/lint.py

@@ -6,6 +6,12 @@ Usage:
    $ python scripts/lint.py rules/
 
 Copyright (C) 2020 FireEye, 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 os
 import sys
@@ -13,11 +19,10 @@ import string
 import hashlib
 import logging
 import os.path
+import argparse
 import itertools
 import posixpath
 
-import argparse
-
 import capa.main
 import capa.engine
 import capa.features
@@ -44,7 +49,8 @@ class NameCasing(Lint):
 
 class FilenameDoesntMatchRuleName(Lint):
     name = "filename doesn't match the rule name"
-    recommendation = 'Rename rule file to match the rule name, expected: "{:s}", found: "{:s}"'
+    recommendation = "Rename rule file to match the rule name"
+    recommendation_template = 'Rename rule file to match the rule name, expected: "{:s}", found: "{:s}"'
 
     def check_rule(self, ctx, rule):
         expected = rule.name
@@ -54,11 +60,12 @@ class FilenameDoesntMatchRuleName(Lint):
         expected = expected.replace(")", "")
         expected = expected.replace("+", "")
         expected = expected.replace("/", "")
+        expected = expected.replace(".", "")
         expected = expected + ".yml"
 
         found = os.path.basename(rule.meta["capa/path"])
 
-        self.recommendation = self.recommendation.format(expected, found)
+        self.recommendation = self.recommendation_template.format(expected, found)
 
         return expected != found
 
@@ -137,10 +144,12 @@ class MissingExampleOffset(Lint):
 
     def check_rule(self, ctx, rule):
         if rule.meta.get("scope") in ("function", "basic block"):
-            for example in rule.meta.get("examples", []):
-                if example and ":" not in example:
-                    logger.debug("example: %s", example)
-                    return True
+            examples = rule.meta.get("examples")
+            if isinstance(examples, list):
+                for example in examples:
+                    if example and ":" not in example:
+                        logger.debug("example: %s", example)
+                        return True
 
 
 class ExampleFileDNE(Lint):
@@ -171,7 +180,11 @@ class DoesntMatchExample(Lint):
         if not ctx["is_thorough"]:
             return False
 
-        for example in rule.meta.get("examples", []):
+        examples = rule.meta.get("examples", [])
+        if not examples:
+            return False
+
+        for example in examples:
             example_id = example.partition(":")[0]
             try:
                 path = ctx["samples"][example_id]
@@ -193,7 +206,8 @@ class DoesntMatchExample(Lint):
 
 class UnusualMetaField(Lint):
     name = "unusual meta field"
-    recommendation = 'Remove the meta field: "{:s}"'
+    recommendation = "Remove the meta field"
+    recommendation_template = 'Remove the meta field: "{:s}"'
 
     def check_rule(self, ctx, rule):
         for key in rule.meta.keys():
@@ -201,7 +215,7 @@ class UnusualMetaField(Lint):
                 continue
             if key in capa.rules.HIDDEN_META_KEYS:
                 continue
-            self.recommendation = self.recommendation.format(key)
+            self.recommendation = self.recommendation_template.format(key)
             return True
 
         return False
@@ -245,18 +259,19 @@ class FeatureStringTooShort(Lint):
         return False
 
 
-class FeatureNegativeNumberOrOffset(Lint):
+class FeatureNegativeNumber(Lint):
     name = "feature value is negative"
-    recommendation = (
-        "capa treats all numbers as unsigned values; you may specify the number's two's complement "
+    recommendation = "specify the number's two's complement representation"
+    recommendation_template = (
+        "capa treats number features as unsigned values; you may specify the number's two's complement "
         'representation; will not match on "{:d}"'
     )
 
     def check_features(self, ctx, features):
         for feature in features:
-            if isinstance(feature, (capa.features.insn.Number, capa.features.insn.Offset)):
+            if isinstance(feature, (capa.features.insn.Number,)):
                 if feature.value < 0:
-                    self.recommendation = self.recommendation.format(feature.value)
+                    self.recommendation = self.recommendation_template.format(feature.value)
                     return True
         return False
 
@@ -312,7 +327,7 @@ def lint_meta(ctx, rule):
 
 FEATURE_LINTS = (
     FeatureStringTooShort(),
-    FeatureNegativeNumberOrOffset(),
+    FeatureNegativeNumber(),
 )
 
 
@@ -384,7 +399,11 @@ def lint_rule(ctx, rule):
         print("")
         print(
             "%s%s %s"
-            % ("    (nursery) " if is_nursery_rule(rule) else "", rule.name, ("(%s)" % category) if category else "",)
+            % (
+                "    (nursery) " if is_nursery_rule(rule) else "",
+                rule.name,
+                ("(%s)" % category) if category else "",
+            )
         )
 
         level = "WARN" if is_nursery_rule(rule) else "FAIL"
@@ -392,9 +411,19 @@ def lint_rule(ctx, rule):
         for violation in violations:
             print(
                 "%s  %s: %s: %s"
-                % ("    " if is_nursery_rule(rule) else "", level, violation.name, violation.recommendation,)
+                % (
+                    "    " if is_nursery_rule(rule) else "",
+                    level,
+                    violation.name,
+                    violation.recommendation,
+                )
             )
 
+    elif len(violations) == 0 and is_nursery_rule(rule):
+        print("")
+        print("%s%s" % ("    (nursery) ", rule.name))
+        print("%s  %s: %s: %s" % ("    ", "WARN", "no violations", "Graduate the rule"))
+
     return len(violations) > 0 and not is_nursery_rule(rule)
 
 
@@ -467,7 +496,9 @@ def main(argv=None):
     parser.add_argument("rules", type=str, help="Path to rules")
     parser.add_argument("--samples", type=str, default=samples_path, help="Path to samples")
     parser.add_argument(
-        "--thorough", action="store_true", help="Enable thorough linting - takes more time, but does a better job",
+        "--thorough",
+        action="store_true",
+        help="Enable thorough linting - takes more time, but does a better job",
     )
     parser.add_argument("-v", "--verbose", action="store_true", help="Enable debug logging")
     parser.add_argument("-q", "--quiet", action="store_true", help="Disable all output but errors")

scripts/migrate-rules.py

@@ -7,15 +7,20 @@ example:
     $ python scripts/migrate-rules.py migration.csv ./rules ./new-rules
 
 Copyright (C) 2020 FireEye, 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 os
 import csv
 import sys
 import logging
 import os.path
-import collections
-
 import argparse
+import collections
 
 import capa.rules
 

scripts/setup-hooks.sh

@@ -1,5 +1,11 @@
 #!/usr/bin/env bash
 # Copyright (C) 2020 FireEye, 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.
 
 set -euo pipefail
 
@@ -7,16 +13,16 @@ GIT_DIR=$(git rev-parse --show-toplevel);
 cd "$GIT_DIR";
 
 # hooks may exist already (e.g. git-lfs configuration)
-# If the `.git/hooks/$arg` file doesn't exist it, initialize with `#!/bin/sh`
+# If the `.git/hooks/$arg` file doesn't exist it, initialize with `#!/usr/bin/env bash`
 # After that append `scripts/hooks/$arg` and ensure they can be run
 create_hook() {
   if [[ ! -e .git/hooks/$1 ]]; then
-    echo "#!/bin/sh" > ".git/hooks/$1";
+    echo "#!/usr/bin/env bash" > ".git/hooks/$1";
   fi
-  cat scripts/hooks/"$1" >> ".git/hooks/$1";
+  echo "scripts/ci.sh ${2:-}" >> ".git/hooks/$1";
   chmod +x .git/hooks/"$1";
 }
 
-printf '\n#### Copying hooks into .git/hooks';
-create_hook 'post-commit';
+printf 'Adding scripts/ci.sh to .git/hooks/';
+create_hook 'pre-commit' 'no_tests';
 create_hook 'pre-push';

scripts/show-capabilities-by-function.py

@@ -41,13 +41,19 @@ Example::
     ...
 
 Copyright (C) 2020 FireEye, 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 os
 import sys
 import logging
+import argparse
 import collections
 
-import argparse
 import colorama
 
 import capa.main
@@ -65,22 +71,22 @@ logger = logging.getLogger("capa.show-capabilities-by-function")
 
 def render_matches_by_function(doc):
     """
-        like:
+    like:
 
-            function at 0x1000321a with 33 features:
-              - get hostname
-              - initialize Winsock library
-            function at 0x10003286 with 63 features:
-              - create thread
-              - terminate thread
-            function at 0x10003415 with 116 features:
-              - write file
-              - send data
-              - link function at runtime
-              - create HTTP request
-              - get common file path
-              - send HTTP request
-              - connect to HTTP server
+        function at 0x1000321a with 33 features:
+          - get hostname
+          - initialize Winsock library
+        function at 0x10003286 with 63 features:
+          - create thread
+          - terminate thread
+        function at 0x10003415 with 116 features:
+          - write file
+          - send data
+          - link function at runtime
+          - create HTTP request
+          - get common file path
+          - send HTTP request
+          - connect to HTTP server
     """
     ostream = rutils.StringIO()
 

scripts/show-features.py

@@ -57,10 +57,15 @@ Example::
     ...
 
 Copyright (C) 2020 FireEye, 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 sys
 import logging
-
 import argparse
 
 import capa.main
@@ -127,7 +132,11 @@ def main(argv=None):
 
             for insn in extractor.get_instructions(f, bb):
                 for feature, va in extractor.extract_insn_features(f, bb, insn):
-                    print("insn: 0x%08x: %s" % (va, feature))
+                    try:
+                        print("insn: 0x%08x: %s" % (va, feature))
+                    except UnicodeEncodeError:
+                        # may be an issue while piping to less and encountering non-ascii characters
+                        continue
 
     return 0
 

setup.py

@@ -1,21 +1,40 @@
 # Copyright (C) 2020 FireEye, 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 os
 import sys
 
 import setuptools
 
-requirements = ["six", "tqdm", "pyyaml", "tabulate", "colorama", "termcolor", "ruamel.yaml"]
+# halo==0.0.30 is the last version to support py2.7
+requirements = [
+    "six",
+    "tqdm",
+    "pyyaml",
+    "tabulate",
+    "colorama",
+    "termcolor",
+    "ruamel.yaml",
+    "wcwidth",
+    "halo==0.0.30",
+    "ida-settings==2.1.0",
+]
 
 if sys.version_info >= (3, 0):
     # py3
     requirements.append("networkx")
 else:
     # py2
-    requirements.append("enum34")
-    requirements.append("vivisect")
+    requirements.append("enum34==1.1.6")  # v1.1.6 is needed by halo 0.0.30 / spinners 0.0.24
+    requirements.append("vivisect==0.1.0")
     requirements.append("viv-utils")
     requirements.append("networkx==2.2")  # v2.2 is last version supported by Python 2.7
+    requirements.append("backports.functools-lru-cache")
 
 # this sets __version__
 # via: http://stackoverflow.com/a/7071358/87207
@@ -25,7 +44,7 @@ with open(os.path.join("capa", "version.py"), "rb") as f:
 
 
 setuptools.setup(
-    name="capa",
+    name="flare-capa",
     version=__version__,
     description="The FLARE team's open-source tool to identify capabilities in executable files.",
     long_description="",
@@ -34,7 +53,11 @@ setuptools.setup(
     url="https://www.github.com/fireeye/capa",
     packages=setuptools.find_packages(exclude=["tests"]),
     package_dir={"capa": "capa"},
-    entry_points={"console_scripts": ["capa=capa.main:main",]},
+    entry_points={
+        "console_scripts": [
+            "capa=capa.main:main",
+        ]
+    },
     include_package_data=True,
     install_requires=requirements,
     extras_require={

tests/fixtures.py

@@ -1,79 +1,528 @@
 # Copyright (C) 2020 FireEye, 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 os
+import sys
 import os.path
+import contextlib
 import collections
 
 import pytest
-import viv_utils
+
+import capa.main
+import capa.features.file
+import capa.features.insn
+import capa.features.basicblock
+from capa.features import ARCH_X32, ARCH_X64
+
+try:
+    from functools import lru_cache
+except ImportError:
+    from backports.functools_lru_cache import lru_cache
+
 
 CD = os.path.dirname(__file__)
 
 
-Sample = collections.namedtuple("Sample", ["vw", "path"])
+@contextlib.contextmanager
+def xfail(condition, reason=None):
+    """
+    context manager that wraps a block that is expected to fail in some cases.
+    when it does fail (and is expected), then mark this as pytest.xfail.
+    if its unexpected, raise an exception, so the test fails.
+
+    example::
+
+        # this test:
+        #  - passes on py3 if foo() works
+        #  - fails  on py3 if foo() fails
+        #  - xfails on py2 if foo() fails
+        #  - fails  on py2 if foo() works
+        with xfail(sys.version_info < (3, 0), reason="py2 doesn't foo"):
+            foo()
+    """
+    try:
+        # do the block
+        yield
+    except:
+        if condition:
+            # we expected the test to fail, so raise and register this via pytest
+            pytest.xfail(reason)
+        else:
+            # we don't expect an exception, so the test should fail
+            raise
+    else:
+        if not condition:
+            # here we expect the block to run successfully,
+            # and we've received no exception,
+            # so this is good
+            pass
+        else:
+            # we expected an exception, but didn't find one. that's an error.
+            raise RuntimeError("expected to fail, but didn't")
+
+
+@lru_cache()
+def get_viv_extractor(path):
+    import capa.features.extractors.viv
+
+    if "raw32" in path:
+        vw = capa.main.get_workspace(path, "sc32", should_save=False)
+    elif "raw64" in path:
+        vw = capa.main.get_workspace(path, "sc64", should_save=False)
+    else:
+        vw = capa.main.get_workspace(path, "auto", should_save=True)
+    return capa.features.extractors.viv.VivisectFeatureExtractor(vw, path)
+
+
+@lru_cache()
+def extract_file_features(extractor):
+    features = collections.defaultdict(set)
+    for feature, va in extractor.extract_file_features():
+        features[feature].add(va)
+    return features
+
+
+# f may not be hashable (e.g. ida func_t) so cannot @lru_cache this
+def extract_function_features(extractor, f):
+    features = collections.defaultdict(set)
+    for bb in extractor.get_basic_blocks(f):
+        for insn in extractor.get_instructions(f, bb):
+            for feature, va in extractor.extract_insn_features(f, bb, insn):
+                features[feature].add(va)
+        for feature, va in extractor.extract_basic_block_features(f, bb):
+            features[feature].add(va)
+    for feature, va in extractor.extract_function_features(f):
+        features[feature].add(va)
+    return features
+
+
+# f may not be hashable (e.g. ida func_t) so cannot @lru_cache this
+def extract_basic_block_features(extractor, f, bb):
+    features = collections.defaultdict(set)
+    for insn in extractor.get_instructions(f, bb):
+        for feature, va in extractor.extract_insn_features(f, bb, insn):
+            features[feature].add(va)
+    for feature, va in extractor.extract_basic_block_features(f, bb):
+        features[feature].add(va)
+    return features
+
+
+def get_data_path_by_name(name):
+    if name == "mimikatz":
+        return os.path.join(CD, "data", "mimikatz.exe_")
+    elif name == "kernel32":
+        return os.path.join(CD, "data", "kernel32.dll_")
+    elif name == "kernel32-64":
+        return os.path.join(CD, "data", "kernel32-64.dll_")
+    elif name == "pma12-04":
+        return os.path.join(CD, "data", "Practical Malware Analysis Lab 12-04.exe_")
+    elif name == "pma16-01":
+        return os.path.join(CD, "data", "Practical Malware Analysis Lab 16-01.exe_")
+    elif name == "pma21-01":
+        return os.path.join(CD, "data", "Practical Malware Analysis Lab 21-01.exe_")
+    elif name == "al-khaser x86":
+        return os.path.join(CD, "data", "al-khaser_x86.exe_")
+    elif name.startswith("39c05"):
+        return os.path.join(CD, "data", "39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.dll_")
+    elif name.startswith("499c2"):
+        return os.path.join(CD, "data", "499c2a85f6e8142c3f48d4251c9c7cd6.raw32")
+    elif name.startswith("9324d"):
+        return os.path.join(CD, "data", "9324d1a8ae37a36ae560c37448c9705a.exe_")
+    elif name.startswith("a1982"):
+        return os.path.join(CD, "data", "a198216798ca38f280dc413f8c57f2c2.exe_")
+    elif name.startswith("a933a"):
+        return os.path.join(CD, "data", "a933a1a402775cfa94b6bee0963f4b46.dll_")
+    elif name.startswith("bfb9b"):
+        return os.path.join(CD, "data", "bfb9b5391a13d0afd787e87ab90f14f5.dll_")
+    elif name.startswith("c9188"):
+        return os.path.join(CD, "data", "c91887d861d9bd4a5872249b641bc9f9.exe_")
+    elif name.startswith("64d9f"):
+        return os.path.join(CD, "data", "64d9f7d96b99467f36e22fada623c3bb.dll_")
+    else:
+        raise ValueError("unexpected sample fixture")
+
+
+def get_sample_md5_by_name(name):
+    """used by IDA tests to ensure the correct IDB is loaded"""
+    if name == "mimikatz":
+        return "5f66b82558ca92e54e77f216ef4c066c"
+    elif name == "kernel32":
+        return "e80758cf485db142fca1ee03a34ead05"
+    elif name == "kernel32-64":
+        return "a8565440629ac87f6fef7d588fe3ff0f"
+    elif name == "pma12-04":
+        return "56bed8249e7c2982a90e54e1e55391a2"
+    elif name == "pma16-01":
+        return "7faafc7e4a5c736ebfee6abbbc812d80"
+    elif name == "pma21-01":
+        return "c8403fb05244e23a7931c766409b5e22"
+    elif name == "al-khaser x86":
+        return "db648cd247281954344f1d810c6fd590"
+    elif name.startswith("39c05"):
+        return "b7841b9d5dc1f511a93cc7576672ec0c"
+    elif name.startswith("499c2"):
+        return "499c2a85f6e8142c3f48d4251c9c7cd6"
+    elif name.startswith("9324d"):
+        return "9324d1a8ae37a36ae560c37448c9705a"
+    elif name.startswith("a1982"):
+        return "a198216798ca38f280dc413f8c57f2c2"
+    elif name.startswith("a933a"):
+        return "a933a1a402775cfa94b6bee0963f4b46"
+    elif name.startswith("bfb9b"):
+        return "bfb9b5391a13d0afd787e87ab90f14f5"
+    elif name.startswith("c9188"):
+        return "c91887d861d9bd4a5872249b641bc9f9"
+    elif name.startswith("64d9f"):
+        return "64d9f7d96b99467f36e22fada623c3bb"
+    else:
+        raise ValueError("unexpected sample fixture")
+
+
+def resolve_sample(sample):
+    return get_data_path_by_name(sample)
 
 
 @pytest.fixture
-def mimikatz():
-    path = os.path.join(CD, "data", "mimikatz.exe_")
-    return Sample(viv_utils.getWorkspace(path), path)
+def sample(request):
+    return resolve_sample(request.param)
+
+
+def get_function(extractor, fva):
+    for f in extractor.get_functions():
+        if f.__int__() == fva:
+            return f
+    raise ValueError("function not found")
+
+
+def get_basic_block(extractor, f, va):
+    for bb in extractor.get_basic_blocks(f):
+        if bb.__int__() == va:
+            return bb
+    raise ValueError("basic block not found")
+
+
+def resolve_scope(scope):
+    if scope == "file":
+
+        def inner(extractor):
+            return extract_file_features(extractor)
+
+        inner.__name__ = scope
+        return inner
+    elif "bb=" in scope:
+        # like `function=0x401000,bb=0x40100A`
+        fspec, _, bbspec = scope.partition(",")
+        fva = int(fspec.partition("=")[2], 0x10)
+        bbva = int(bbspec.partition("=")[2], 0x10)
+
+        def inner(extractor):
+            f = get_function(extractor, fva)
+            bb = get_basic_block(extractor, f, bbva)
+            return extract_basic_block_features(extractor, f, bb)
+
+        inner.__name__ = scope
+        return inner
+    elif scope.startswith("function"):
+        # like `function=0x401000`
+        va = int(scope.partition("=")[2], 0x10)
+
+        def inner(extractor):
+            f = get_function(extractor, va)
+            return extract_function_features(extractor, f)
+
+        inner.__name__ = scope
+        return inner
+    else:
+        raise ValueError("unexpected scope fixture")
 
 
 @pytest.fixture
-def sample_a933a1a402775cfa94b6bee0963f4b46():
-    path = os.path.join(CD, "data", "a933a1a402775cfa94b6bee0963f4b46.dll_")
-    return Sample(viv_utils.getWorkspace(path), path)
+def scope(request):
+    return resolve_scope(request.param)
+
+
+def make_test_id(values):
+    return "-".join(map(str, values))
+
+
+def parametrize(params, values, **kwargs):
+    """
+    extend `pytest.mark.parametrize` to pretty-print features.
+    by default, it renders objects as an opaque value.
+    ref: https://docs.pytest.org/en/2.9.0/example/parametrize.html#different-options-for-test-ids
+    rendered ID might look something like:
+        mimikatz-function=0x403BAC-api(CryptDestroyKey)-True
+    """
+    ids = list(map(make_test_id, values))
+    return pytest.mark.parametrize(params, values, ids=ids, **kwargs)
+
+
+FEATURE_PRESENCE_TESTS = [
+    # file/characteristic("embedded pe")
+    ("pma12-04", "file", capa.features.Characteristic("embedded pe"), True),
+    # file/string
+    ("mimikatz", "file", capa.features.String("SCardControl"), True),
+    ("mimikatz", "file", capa.features.String("SCardTransmit"), True),
+    ("mimikatz", "file", capa.features.String("ACR  > "), True),
+    ("mimikatz", "file", capa.features.String("nope"), False),
+    # file/sections
+    ("mimikatz", "file", capa.features.file.Section(".text"), True),
+    ("mimikatz", "file", capa.features.file.Section(".nope"), False),
+    # IDA doesn't extract unmapped sections by default
+    # ("mimikatz", "file", capa.features.file.Section(".rsrc"), True),
+    # file/exports
+    ("kernel32", "file", capa.features.file.Export("BaseThreadInitThunk"), True),
+    ("kernel32", "file", capa.features.file.Export("lstrlenW"), True),
+    ("kernel32", "file", capa.features.file.Export("nope"), False),
+    # file/imports
+    ("mimikatz", "file", capa.features.file.Import("advapi32.CryptSetHashParam"), True),
+    ("mimikatz", "file", capa.features.file.Import("CryptSetHashParam"), True),
+    ("mimikatz", "file", capa.features.file.Import("kernel32.IsWow64Process"), True),
+    ("mimikatz", "file", capa.features.file.Import("msvcrt.exit"), True),
+    ("mimikatz", "file", capa.features.file.Import("cabinet.#11"), True),
+    ("mimikatz", "file", capa.features.file.Import("#11"), False),
+    ("mimikatz", "file", capa.features.file.Import("#nope"), False),
+    ("mimikatz", "file", capa.features.file.Import("nope"), False),
+    ("mimikatz", "file", capa.features.file.Import("advapi32.CryptAcquireContextW"), True),
+    ("mimikatz", "file", capa.features.file.Import("advapi32.CryptAcquireContext"), True),
+    ("mimikatz", "file", capa.features.file.Import("CryptAcquireContextW"), True),
+    ("mimikatz", "file", capa.features.file.Import("CryptAcquireContext"), True),
+    # function/characteristic(loop)
+    ("mimikatz", "function=0x401517", capa.features.Characteristic("loop"), True),
+    ("mimikatz", "function=0x401000", capa.features.Characteristic("loop"), False),
+    # bb/characteristic(tight loop)
+    ("mimikatz", "function=0x402EC4", capa.features.Characteristic("tight loop"), True),
+    ("mimikatz", "function=0x401000", capa.features.Characteristic("tight loop"), False),
+    # bb/characteristic(stack string)
+    ("mimikatz", "function=0x4556E5", capa.features.Characteristic("stack string"), True),
+    ("mimikatz", "function=0x401000", capa.features.Characteristic("stack string"), False),
+    # bb/characteristic(tight loop)
+    ("mimikatz", "function=0x402EC4,bb=0x402F8E", capa.features.Characteristic("tight loop"), True),
+    ("mimikatz", "function=0x401000,bb=0x401000", capa.features.Characteristic("tight loop"), False),
+    # insn/mnemonic
+    ("mimikatz", "function=0x40105D", capa.features.insn.Mnemonic("push"), True),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Mnemonic("movzx"), True),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Mnemonic("xor"), True),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Mnemonic("in"), False),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Mnemonic("out"), False),
+    # insn/number
+    ("mimikatz", "function=0x40105D", capa.features.insn.Number(0xFF), True),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Number(0x3136B0), True),
+    # insn/number: stack adjustments
+    ("mimikatz", "function=0x40105D", capa.features.insn.Number(0xC), False),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Number(0x10), False),
+    # insn/number: arch flavors
+    ("mimikatz", "function=0x40105D", capa.features.insn.Number(0xFF), True),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Number(0xFF, arch=ARCH_X32), True),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Number(0xFF, arch=ARCH_X64), False),
+    # insn/offset
+    ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x0), True),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x4), True),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0xC), True),
+    # insn/offset, issue #276
+    ("64d9f", "function=0x10001510,bb=0x100015B0", capa.features.insn.Offset(0x4000), True),
+    # insn/offset: stack references
+    ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x8), False),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x10), False),
+    # insn/offset: negative
+    ("mimikatz", "function=0x4011FB", capa.features.insn.Offset(-0x1), True),
+    ("mimikatz", "function=0x4011FB", capa.features.insn.Offset(-0x2), True),
+    # insn/offset: arch flavors
+    ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x0), True),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x0, arch=ARCH_X32), True),
+    ("mimikatz", "function=0x40105D", capa.features.insn.Offset(0x0, arch=ARCH_X64), False),
+    # insn/api
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.CryptAcquireContextW"), True),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.CryptAcquireContext"), True),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.CryptGenKey"), True),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.CryptImportKey"), True),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.CryptDestroyKey"), True),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("CryptAcquireContextW"), True),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("CryptAcquireContext"), True),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("CryptGenKey"), True),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("CryptImportKey"), True),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("CryptDestroyKey"), True),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("Nope"), False),
+    ("mimikatz", "function=0x403BAC", capa.features.insn.API("advapi32.Nope"), False),
+    # insn/api: thunk
+    ("mimikatz", "function=0x4556E5", capa.features.insn.API("advapi32.LsaQueryInformationPolicy"), True),
+    ("mimikatz", "function=0x4556E5", capa.features.insn.API("LsaQueryInformationPolicy"), True),
+    # insn/api: x64
+    (
+        "kernel32-64",
+        "function=0x180001010",
+        capa.features.insn.API("RtlVirtualUnwind"),
+        True,
+    ),
+    ("kernel32-64", "function=0x180001010", capa.features.insn.API("RtlVirtualUnwind"), True),
+    # insn/api: x64 thunk
+    (
+        "kernel32-64",
+        "function=0x1800202B0",
+        capa.features.insn.API("RtlCaptureContext"),
+        True,
+    ),
+    ("kernel32-64", "function=0x1800202B0", capa.features.insn.API("RtlCaptureContext"), True),
+    # insn/api: resolve indirect calls
+    ("c91887...", "function=0x401A77", capa.features.insn.API("kernel32.CreatePipe"), True),
+    ("c91887...", "function=0x401A77", capa.features.insn.API("kernel32.SetHandleInformation"), True),
+    ("c91887...", "function=0x401A77", capa.features.insn.API("kernel32.CloseHandle"), True),
+    ("c91887...", "function=0x401A77", capa.features.insn.API("kernel32.WriteFile"), True),
+    # insn/string
+    ("mimikatz", "function=0x40105D", capa.features.String("SCardControl"), True),
+    ("mimikatz", "function=0x40105D", capa.features.String("SCardTransmit"), True),
+    ("mimikatz", "function=0x40105D", capa.features.String("ACR  > "), True),
+    ("mimikatz", "function=0x40105D", capa.features.String("nope"), False),
+    # insn/regex, issue #262
+    ("pma16-01", "function=0x4021B0", capa.features.Regex("HTTP/1.0"), True),
+    ("pma16-01", "function=0x4021B0", capa.features.Regex("www.practicalmalwareanalysis.com"), False),
+    # insn/string, pointer to string
+    ("mimikatz", "function=0x44EDEF", capa.features.String("INPUTEVENT"), True),
+    # insn/bytes
+    ("mimikatz", "function=0x40105D", capa.features.Bytes("SCardControl".encode("utf-16le")), True),
+    ("mimikatz", "function=0x40105D", capa.features.Bytes("SCardTransmit".encode("utf-16le")), True),
+    ("mimikatz", "function=0x40105D", capa.features.Bytes("ACR  > ".encode("utf-16le")), True),
+    ("mimikatz", "function=0x40105D", capa.features.Bytes("nope".encode("ascii")), False),
+    # insn/bytes, pointer to bytes
+    ("mimikatz", "function=0x44EDEF", capa.features.Bytes("INPUTEVENT".encode("utf-16le")), True),
+    # insn/characteristic(nzxor)
+    ("mimikatz", "function=0x410DFC", capa.features.Characteristic("nzxor"), True),
+    ("mimikatz", "function=0x40105D", capa.features.Characteristic("nzxor"), False),
+    # insn/characteristic(nzxor): no security cookies
+    ("mimikatz", "function=0x46D534", capa.features.Characteristic("nzxor"), False),
+    # insn/characteristic(peb access)
+    ("kernel32-64", "function=0x1800017D0", capa.features.Characteristic("peb access"), True),
+    ("mimikatz", "function=0x4556E5", capa.features.Characteristic("peb access"), False),
+    # insn/characteristic(gs access)
+    ("kernel32-64", "function=0x180001068", capa.features.Characteristic("gs access"), True),
+    ("mimikatz", "function=0x4556E5", capa.features.Characteristic("gs access"), False),
+    # insn/characteristic(cross section flow)
+    ("a1982...", "function=0x4014D0", capa.features.Characteristic("cross section flow"), True),
+    # insn/characteristic(cross section flow): imports don't count
+    ("kernel32-64", "function=0x180001068", capa.features.Characteristic("cross section flow"), False),
+    ("mimikatz", "function=0x4556E5", capa.features.Characteristic("cross section flow"), False),
+    # insn/characteristic(recursive call)
+    ("39c05...", "function=0x10003100", capa.features.Characteristic("recursive call"), True),
+    ("mimikatz", "function=0x4556E5", capa.features.Characteristic("recursive call"), False),
+    # insn/characteristic(indirect call)
+    ("mimikatz", "function=0x4175FF", capa.features.Characteristic("indirect call"), True),
+    ("mimikatz", "function=0x4556E5", capa.features.Characteristic("indirect call"), False),
+    # insn/characteristic(calls from)
+    ("mimikatz", "function=0x4556E5", capa.features.Characteristic("calls from"), True),
+    ("mimikatz", "function=0x4702FD", capa.features.Characteristic("calls from"), False),
+    # function/characteristic(calls to)
+    ("mimikatz", "function=0x40105D", capa.features.Characteristic("calls to"), True),
+    ("mimikatz", "function=0x4556E5", capa.features.Characteristic("calls to"), False),
+]
+
+FEATURE_COUNT_TESTS = [
+    ("mimikatz", "function=0x40E5C2", capa.features.basicblock.BasicBlock(), 7),
+    ("mimikatz", "function=0x4702FD", capa.features.Characteristic("calls from"), 0),
+    ("mimikatz", "function=0x40E5C2", capa.features.Characteristic("calls from"), 3),
+    ("mimikatz", "function=0x4556E5", capa.features.Characteristic("calls to"), 0),
+    ("mimikatz", "function=0x40B1F1", capa.features.Characteristic("calls to"), 3),
+]
+
+
+def do_test_feature_presence(get_extractor, sample, scope, feature, expected):
+    extractor = get_extractor(sample)
+    features = scope(extractor)
+    if expected:
+        msg = "%s should be found in %s" % (str(feature), scope.__name__)
+    else:
+        msg = "%s should not be found in %s" % (str(feature), scope.__name__)
+    assert feature.evaluate(features) == expected, msg
+
+
+def do_test_feature_count(get_extractor, sample, scope, feature, expected):
+    extractor = get_extractor(sample)
+    features = scope(extractor)
+    msg = "%s should be found %d times in %s, found: %d" % (
+        str(feature),
+        expected,
+        scope.__name__,
+        len(features[feature]),
+    )
+    assert len(features[feature]) == expected, msg
+
+
+def get_extractor(path):
+    if sys.version_info >= (3, 0):
+        raise RuntimeError("no supported py3 backends yet")
+    else:
+        extractor = get_viv_extractor(path)
+
+    # overload the extractor so that the fixture exposes `extractor.path`
+    setattr(extractor, "path", path)
+    return extractor
 
 
 @pytest.fixture
-def kernel32():
-    path = os.path.join(CD, "data", "kernel32.dll_")
-    return Sample(viv_utils.getWorkspace(path), path)
+def mimikatz_extractor():
+    return get_extractor(get_data_path_by_name("mimikatz"))
 
 
 @pytest.fixture
-def sample_a198216798ca38f280dc413f8c57f2c2():
-    path = os.path.join(CD, "data", "a198216798ca38f280dc413f8c57f2c2.exe_")
-    return Sample(viv_utils.getWorkspace(path), path)
+def a933a_extractor():
+    return get_extractor(get_data_path_by_name("a933a..."))
 
 
 @pytest.fixture
-def sample_9324d1a8ae37a36ae560c37448c9705a():
-    path = os.path.join(CD, "data", "9324d1a8ae37a36ae560c37448c9705a.exe_")
-    return Sample(viv_utils.getWorkspace(path), path)
+def kernel32_extractor():
+    return get_extractor(get_data_path_by_name("kernel32"))
 
 
 @pytest.fixture
-def pma_lab_12_04():
-    path = os.path.join(CD, "data", "Practical Malware Analysis Lab 12-04.exe_")
-    return Sample(viv_utils.getWorkspace(path), path)
+def a1982_extractor():
+    return get_extractor(get_data_path_by_name("a1982..."))
 
 
 @pytest.fixture
-def sample_bfb9b5391a13d0afd787e87ab90f14f5():
-    path = os.path.join(CD, "data", "bfb9b5391a13d0afd787e87ab90f14f5.dll_")
-    return Sample(viv_utils.getWorkspace(path), path)
+def z9324d_extractor():
+    return get_extractor(get_data_path_by_name("9324d..."))
 
 
 @pytest.fixture
-def sample_lab21_01():
-    path = os.path.join(CD, "data", "Practical Malware Analysis Lab 21-01.exe_")
-    return Sample(viv_utils.getWorkspace(path), path)
+def pma12_04_extractor():
+    return get_extractor(get_data_path_by_name("pma12-04"))
 
 
 @pytest.fixture
-def sample_c91887d861d9bd4a5872249b641bc9f9():
-    path = os.path.join(CD, "data", "c91887d861d9bd4a5872249b641bc9f9.exe_")
-    return Sample(viv_utils.getWorkspace(path), path)
+def pma16_01_extractor():
+    return get_extractor(get_data_path_by_name("pma16-01"))
 
 
 @pytest.fixture
-def sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41():
-    path = os.path.join(CD, "data", "39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.dll_",)
-    return Sample(viv_utils.getWorkspace(path), path)
+def bfb9b_extractor():
+    return get_extractor(get_data_path_by_name("bfb9b..."))
 
 
 @pytest.fixture
-def sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32():
-    path = os.path.join(CD, "data", "499c2a85f6e8142c3f48d4251c9c7cd6.raw32")
-    return Sample(viv_utils.getShellcodeWorkspace(path), path)
+def pma21_01_extractor():
+    return get_extractor(get_data_path_by_name("pma21-01"))
+
+
+@pytest.fixture
+def c9188_extractor():
+    return get_extractor(get_data_path_by_name("c9188..."))
+
+
+@pytest.fixture
+def z39c05_extractor():
+    return get_extractor(get_data_path_by_name("39c05..."))
+
+
+@pytest.fixture
+def z499c2_extractor():
+    return get_extractor(get_data_path_by_name("499c2..."))
+
+
+@pytest.fixture
+def al_khaser_x86_extractor():
+    return get_extractor(get_data_path_by_name("al-khaser x86"))

tests/test_engine.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 textwrap
 
@@ -16,21 +22,21 @@ def test_number():
 
 
 def test_and():
-    assert And(Number(1)).evaluate({Number(0): {1}}) == False
-    assert And(Number(1)).evaluate({Number(1): {1}}) == True
-    assert And(Number(1), Number(2)).evaluate({Number(0): {1}}) == False
-    assert And(Number(1), Number(2)).evaluate({Number(1): {1}}) == False
-    assert And(Number(1), Number(2)).evaluate({Number(2): {1}}) == False
-    assert And(Number(1), Number(2)).evaluate({Number(1): {1}, Number(2): {2}}) == True
+    assert And([Number(1)]).evaluate({Number(0): {1}}) == False
+    assert And([Number(1)]).evaluate({Number(1): {1}}) == True
+    assert And([Number(1), Number(2)]).evaluate({Number(0): {1}}) == False
+    assert And([Number(1), Number(2)]).evaluate({Number(1): {1}}) == False
+    assert And([Number(1), Number(2)]).evaluate({Number(2): {1}}) == False
+    assert And([Number(1), Number(2)]).evaluate({Number(1): {1}, Number(2): {2}}) == True
 
 
 def test_or():
-    assert Or(Number(1)).evaluate({Number(0): {1}}) == False
-    assert Or(Number(1)).evaluate({Number(1): {1}}) == True
-    assert Or(Number(1), Number(2)).evaluate({Number(0): {1}}) == False
-    assert Or(Number(1), Number(2)).evaluate({Number(1): {1}}) == True
-    assert Or(Number(1), Number(2)).evaluate({Number(2): {1}}) == True
-    assert Or(Number(1), Number(2)).evaluate({Number(1): {1}, Number(2): {2}}) == True
+    assert Or([Number(1)]).evaluate({Number(0): {1}}) == False
+    assert Or([Number(1)]).evaluate({Number(1): {1}}) == True
+    assert Or([Number(1), Number(2)]).evaluate({Number(0): {1}}) == False
+    assert Or([Number(1), Number(2)]).evaluate({Number(1): {1}}) == True
+    assert Or([Number(1), Number(2)]).evaluate({Number(2): {1}}) == True
+    assert Or([Number(1), Number(2)]).evaluate({Number(1): {1}, Number(2): {2}}) == True
 
 
 def test_not():
@@ -39,32 +45,38 @@ def test_not():
 
 
 def test_some():
-    assert Some(0, Number(1)).evaluate({Number(0): {1}}) == True
-    assert Some(1, Number(1)).evaluate({Number(0): {1}}) == False
+    assert Some(0, [Number(1)]).evaluate({Number(0): {1}}) == True
+    assert Some(1, [Number(1)]).evaluate({Number(0): {1}}) == False
 
-    assert Some(2, Number(1), Number(2), Number(3)).evaluate({Number(0): {1}}) == False
-    assert Some(2, Number(1), Number(2), Number(3)).evaluate({Number(0): {1}, Number(1): {1}}) == False
-    assert Some(2, Number(1), Number(2), Number(3)).evaluate({Number(0): {1}, Number(1): {1}, Number(2): {1}}) == True
+    assert Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {1}}) == False
+    assert Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {1}, Number(1): {1}}) == False
+    assert Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {1}, Number(1): {1}, Number(2): {1}}) == True
     assert (
-        Some(2, Number(1), Number(2), Number(3)).evaluate(
+        Some(2, [Number(1), Number(2), Number(3)]).evaluate(
             {Number(0): {1}, Number(1): {1}, Number(2): {1}, Number(3): {1}}
         )
         == True
     )
     assert (
-        Some(2, Number(1), Number(2), Number(3)).evaluate(
-            {Number(0): {1}, Number(1): {1}, Number(2): {1}, Number(3): {1}, Number(4): {1},}
+        Some(2, [Number(1), Number(2), Number(3)]).evaluate(
+            {
+                Number(0): {1},
+                Number(1): {1},
+                Number(2): {1},
+                Number(3): {1},
+                Number(4): {1},
+            }
         )
         == True
     )
 
 
 def test_complex():
-    assert True == Or(And(Number(1), Number(2)), Or(Number(3), Some(2, Number(4), Number(5), Number(6))),).evaluate(
-        {Number(5): {1}, Number(6): {1}, Number(7): {1}, Number(8): {1}}
-    )
+    assert True == Or(
+        [And([Number(1), Number(2)]), Or([Number(3), Some(2, [Number(4), Number(5), Number(6)])])]
+    ).evaluate({Number(5): {1}, Number(6): {1}, Number(7): {1}, Number(8): {1}})
 
-    assert False == Or(And(Number(1), Number(2)), Or(Number(3), Some(2, Number(4), Number(5)))).evaluate(
+    assert False == Or([And([Number(1), Number(2)]), Or([Number(3), Some(2, [Number(4), Number(5)])])]).evaluate(
         {Number(5): {1}, Number(6): {1}, Number(7): {1}, Number(8): {1}}
     )
 
@@ -252,7 +264,9 @@ def test_match_matched_rules():
     ]
 
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(rules), {capa.features.insn.Number(100): {1}}, 0x0,
+        capa.engine.topologically_order_rules(rules),
+        {capa.features.insn.Number(100): {1}},
+        0x0,
     )
     assert capa.features.MatchedRule("test rule1") in features
     assert capa.features.MatchedRule("test rule2") in features
@@ -260,7 +274,9 @@ def test_match_matched_rules():
     # the ordering of the rules must not matter,
     # the engine should match rules in an appropriate order.
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(reversed(rules)), {capa.features.insn.Number(100): {1}}, 0x0,
+        capa.engine.topologically_order_rules(reversed(rules)),
+        {capa.features.insn.Number(100): {1}},
+        0x0,
     )
     assert capa.features.MatchedRule("test rule1") in features
     assert capa.features.MatchedRule("test rule2") in features
@@ -306,22 +322,30 @@ def test_regex():
         ),
     ]
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(rules), {capa.features.insn.Number(100): {1}}, 0x0,
+        capa.engine.topologically_order_rules(rules),
+        {capa.features.insn.Number(100): {1}},
+        0x0,
     )
     assert capa.features.MatchedRule("test rule") not in features
 
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(rules), {capa.features.String("aaaa"): {1}}, 0x0,
+        capa.engine.topologically_order_rules(rules),
+        {capa.features.String("aaaa"): {1}},
+        0x0,
     )
     assert capa.features.MatchedRule("test rule") not in features
 
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(rules), {capa.features.String("aBBBBa"): {1}}, 0x0,
+        capa.engine.topologically_order_rules(rules),
+        {capa.features.String("aBBBBa"): {1}},
+        0x0,
     )
     assert capa.features.MatchedRule("test rule") not in features
 
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(rules), {capa.features.String("abbbba"): {1}}, 0x0,
+        capa.engine.topologically_order_rules(rules),
+        {capa.features.String("abbbba"): {1}},
+        0x0,
     )
     assert capa.features.MatchedRule("test rule") in features
     assert capa.features.MatchedRule("rule with implied wildcards") in features
@@ -344,7 +368,9 @@ def test_regex_ignorecase():
         ),
     ]
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(rules), {capa.features.String("aBBBBa"): {1}}, 0x0,
+        capa.engine.topologically_order_rules(rules),
+        {capa.features.String("aBBBBa"): {1}},
+        0x0,
     )
     assert capa.features.MatchedRule("test rule") in features
 
@@ -423,7 +449,9 @@ def test_match_namespace():
     ]
 
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(rules), {capa.features.insn.API("CreateFile"): {1}}, 0x0,
+        capa.engine.topologically_order_rules(rules),
+        {capa.features.insn.API("CreateFile"): {1}},
+        0x0,
     )
     assert "CreateFile API" in matches
     assert "file-create" in matches
@@ -433,8 +461,22 @@ def test_match_namespace():
     assert capa.features.MatchedRule("file/create/CreateFile") in features
 
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(rules), {capa.features.insn.API("WriteFile"): {1}}, 0x0,
+        capa.engine.topologically_order_rules(rules),
+        {capa.features.insn.API("WriteFile"): {1}},
+        0x0,
     )
     assert "WriteFile API" in matches
     assert "file-create" not in matches
     assert "filesystem-any" in matches
+
+
+def test_render_number():
+    assert str(capa.features.insn.Number(1)) == "number(0x1)"
+    assert str(capa.features.insn.Number(1, arch=ARCH_X32)) == "number/x32(0x1)"
+    assert str(capa.features.insn.Number(1, arch=ARCH_X64)) == "number/x64(0x1)"
+
+
+def test_render_offset():
+    assert str(capa.features.insn.Offset(1)) == "offset(0x1)"
+    assert str(capa.features.insn.Offset(1, arch=ARCH_X32)) == "offset/x32(0x1)"
+    assert str(capa.features.insn.Offset(1, arch=ARCH_X64)) == "offset/x64(0x1)"

tests/test_fmt.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 textwrap
 
@@ -86,6 +92,8 @@ def test_rule_reformat_order():
 
 
 def test_rule_reformat_meta_update():
+    # test updating the rule content after parsing
+
     rule = textwrap.dedent(
         """
         rule:
@@ -106,3 +114,24 @@ def test_rule_reformat_meta_update():
     rule = capa.rules.Rule.from_yaml(rule)
     rule.name = "test rule"
     assert rule.to_yaml() == EXPECTED
+
+
+def test_rule_reformat_string_description():
+    # the `description` should be aligned with the preceding feature name.
+    # see #263
+    src = textwrap.dedent(
+        """
+        rule:
+          meta:
+            name: test rule
+            author: user@domain.com
+            scope: function
+          features:
+            - and:
+              - string: foo
+                description: bar
+        """
+    ).lstrip()
+
+    rule = capa.rules.Rule.from_yaml(src)
+    assert rule.to_yaml() == src

tests/test_freeze.py

@@ -1,25 +1,39 @@
 # Copyright (C) 2020 FireEye, 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 sys
 import textwrap
 
+import pytest
+from fixtures import *
+
 import capa.main
 import capa.helpers
 import capa.features
 import capa.features.insn
 import capa.features.freeze
 import capa.features.extractors
-from fixtures import *
 
 EXTRACTOR = capa.features.extractors.NullFeatureExtractor(
     {
         "base address": 0x401000,
-        "file features": [(0x402345, capa.features.Characteristic("embedded pe")),],
+        "file features": [
+            (0x402345, capa.features.Characteristic("embedded pe")),
+        ],
         "functions": {
             0x401000: {
-                "features": [(0x401000, capa.features.Characteristic("switch")),],
+                "features": [
+                    (0x401000, capa.features.Characteristic("indirect call")),
+                ],
                 "basic blocks": {
                     0x401000: {
-                        "features": [(0x401000, capa.features.Characteristic("tight loop")),],
+                        "features": [
+                            (0x401000, capa.features.Characteristic("tight loop")),
+                        ],
                         "instructions": {
                             0x401000: {
                                 "features": [
@@ -27,7 +41,11 @@ EXTRACTOR = capa.features.extractors.NullFeatureExtractor(
                                     (0x401000, capa.features.Characteristic("nzxor")),
                                 ],
                             },
-                            0x401002: {"features": [(0x401002, capa.features.insn.Mnemonic("mov")),],},
+                            0x401002: {
+                                "features": [
+                                    (0x401002, capa.features.insn.Mnemonic("mov")),
+                                ],
+                            },
                         },
                     },
                 },
@@ -97,17 +115,14 @@ def compare_extractors_viv_null(viv_ext, null_ext):
       viv_ext (capa.features.extractors.viv.VivisectFeatureExtractor)
       null_ext (capa.features.extractors.NullFeatureExtractor)
     """
-
-    # TODO: ordering of these things probably doesn't work yet
-
     assert list(viv_ext.extract_file_features()) == list(null_ext.extract_file_features())
-    assert to_int(list(viv_ext.get_functions())) == list(null_ext.get_functions())
+    assert list(map(to_int, viv_ext.get_functions())) == list(null_ext.get_functions())
     for f in viv_ext.get_functions():
-        assert to_int(list(viv_ext.get_basic_blocks(f))) == list(null_ext.get_basic_blocks(to_int(f)))
+        assert list(map(to_int, viv_ext.get_basic_blocks(f))) == list(null_ext.get_basic_blocks(to_int(f)))
         assert list(viv_ext.extract_function_features(f)) == list(null_ext.extract_function_features(to_int(f)))
 
         for bb in viv_ext.get_basic_blocks(f):
-            assert to_int(list(viv_ext.get_instructions(f, bb))) == list(
+            assert list(map(to_int, viv_ext.get_instructions(f, bb))) == list(
                 null_ext.get_instructions(to_int(f), to_int(bb))
             )
             assert list(viv_ext.extract_basic_block_features(f, bb)) == list(
@@ -122,10 +137,7 @@ def compare_extractors_viv_null(viv_ext, null_ext):
 
 def to_int(o):
     """helper to get int value of extractor items"""
-    if isinstance(o, list):
-        return map(lambda x: capa.helpers.oint(x), o)
-    else:
-        return capa.helpers.oint(o)
+    return capa.helpers.oint(o)
 
 
 def test_freeze_s_roundtrip():
@@ -162,18 +174,22 @@ def test_serialize_features():
     roundtrip_feature(capa.features.file.Import("#11"))
 
 
-def test_freeze_sample(tmpdir, sample_9324d1a8ae37a36ae560c37448c9705a):
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_freeze_sample(tmpdir, z9324d_extractor):
     # tmpdir fixture handles cleanup
     o = tmpdir.mkdir("capa").join("test.frz").strpath
-    assert capa.features.freeze.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, o, "-v"]) == 0
+    path = z9324d_extractor.path
+    assert capa.features.freeze.main([path, o, "-v"]) == 0
 
 
-def test_freeze_load_sample(tmpdir, sample_9324d1a8ae37a36ae560c37448c9705a):
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_freeze_load_sample(tmpdir, z9324d_extractor):
     o = tmpdir.mkdir("capa").join("test.frz")
-    viv_extractor = capa.features.extractors.viv.VivisectFeatureExtractor(
-        sample_9324d1a8ae37a36ae560c37448c9705a.vw, sample_9324d1a8ae37a36ae560c37448c9705a.path,
-    )
+
     with open(o.strpath, "wb") as f:
-        f.write(capa.features.freeze.dump(viv_extractor))
-    null_extractor = capa.features.freeze.load(o.open("rb").read())
-    compare_extractors_viv_null(viv_extractor, null_extractor)
+        f.write(capa.features.freeze.dump(z9324d_extractor))
+
+    with open(o.strpath, "rb") as f:
+        null_extractor = capa.features.freeze.load(f.read())
+
+    compare_extractors_viv_null(z9324d_extractor, null_extractor)

tests/test_helpers.py

@@ -1,4 +1,10 @@
 # Copyright (C) 2020 FireEye, 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 codecs
 

tests/test_ida_features.py

@@ -0,0 +1,104 @@
+# run this script from within IDA with ./tests/data/mimikatz.exe open
+import sys
+import logging
+import os.path
+import binascii
+import traceback
+
+import pytest
+
+try:
+    sys.path.append(os.path.dirname(__file__))
+    from fixtures import *
+finally:
+    sys.path.pop()
+
+
+logger = logging.getLogger("test_ida_features")
+
+
+def check_input_file(wanted):
+    import idautils
+
+    # some versions (7.4) of IDA return a truncated version of the MD5.
+    # https://github.com/idapython/bin/issues/11
+    try:
+        found = idautils.GetInputFileMD5()[:31].decode("ascii").lower()
+    except UnicodeDecodeError:
+        # in IDA 7.5 or so, GetInputFileMD5 started returning raw binary
+        # rather than the hex digest
+        found = binascii.hexlify(idautils.GetInputFileMD5()[:15]).decode("ascii").lower()
+
+    if not wanted.startswith(found):
+        raise RuntimeError("please run the tests against sample with MD5: `%s`" % (wanted))
+
+
+def get_ida_extractor(_path):
+    check_input_file("5f66b82558ca92e54e77f216ef4c066c")
+
+    # have to import import this inline so pytest doesn't bail outside of IDA
+    import capa.features.extractors.ida
+
+    return capa.features.extractors.ida.IdaFeatureExtractor()
+
+
+@pytest.mark.skip(reason="IDA Pro tests must be run within IDA")
+def test_ida_features():
+    for (sample, scope, feature, expected) in FEATURE_PRESENCE_TESTS:
+        id = make_test_id((sample, scope, feature, expected))
+
+        try:
+            check_input_file(get_sample_md5_by_name(sample))
+        except RuntimeError:
+            print("SKIP %s" % (id))
+            continue
+
+        scope = resolve_scope(scope)
+        sample = resolve_sample(sample)
+
+        try:
+            do_test_feature_presence(get_ida_extractor, sample, scope, feature, expected)
+        except Exception as e:
+            print("FAIL %s" % (id))
+            traceback.print_exc()
+        else:
+            print("OK   %s" % (id))
+
+
+@pytest.mark.skip(reason="IDA Pro tests must be run within IDA")
+def test_ida_feature_counts():
+    for (sample, scope, feature, expected) in FEATURE_COUNT_TESTS:
+        id = make_test_id((sample, scope, feature, expected))
+
+        try:
+            check_input_file(get_sample_md5_by_name(sample))
+        except RuntimeError:
+            print("SKIP %s" % (id))
+            continue
+
+        scope = resolve_scope(scope)
+        sample = resolve_sample(sample)
+
+        try:
+            do_test_feature_count(get_ida_extractor, sample, scope, feature, expected)
+        except Exception as e:
+            print("FAIL %s" % (id))
+            traceback.print_exc()
+        else:
+            print("OK   %s" % (id))
+
+
+if __name__ == "__main__":
+    print("-" * 80)
+
+    # invoke all functions in this module that start with `test_`
+    for name in dir(sys.modules[__name__]):
+        if not name.startswith("test_"):
+            continue
+
+        test = getattr(sys.modules[__name__], name)
+        logger.debug("invoking test: %s", name)
+        sys.stderr.flush()
+        test()
+
+    print("DONE")

tests/test_main.py

@@ -1,25 +1,35 @@
 # Copyright (C) 2020 FireEye, 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 sys
 import textwrap
 
+import pytest
+from fixtures import *
+
 import capa.main
 import capa.rules
 import capa.engine
 import capa.features
-import capa.features.extractors.viv
-from fixtures import *
 from capa.engine import *
 
 
-def test_main(sample_9324d1a8ae37a36ae560c37448c9705a):
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_main(z9324d_extractor):
     # tests rules can be loaded successfully and all output modes
-    assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, "-vv"]) == 0
-    assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, "-v"]) == 0
-    assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, "-j"]) == 0
-    assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path]) == 0
+    path = z9324d_extractor.path
+    assert capa.main.main([path, "-vv"]) == 0
+    assert capa.main.main([path, "-v"]) == 0
+    assert capa.main.main([path, "-j"]) == 0
+    assert capa.main.main([path]) == 0
 
 
-def test_main_single_rule(sample_9324d1a8ae37a36ae560c37448c9705a, tmpdir):
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_main_single_rule(z9324d_extractor, tmpdir):
     # tests a single rule can be loaded successfully
     RULE_CONTENT = textwrap.dedent(
         """
@@ -31,16 +41,29 @@ def test_main_single_rule(sample_9324d1a8ae37a36ae560c37448c9705a, tmpdir):
               - string: test
         """
     )
+    path = z9324d_extractor.path
     rule_file = tmpdir.mkdir("capa").join("rule.yml")
     rule_file.write(RULE_CONTENT)
-    assert capa.main.main([sample_9324d1a8ae37a36ae560c37448c9705a.path, "-v", "-r", rule_file.strpath,]) == 0
+    assert (
+        capa.main.main(
+            [
+                path,
+                "-v",
+                "-r",
+                rule_file.strpath,
+            ]
+        )
+        == 0
+    )
 
 
-def test_main_shellcode(sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32):
-    assert capa.main.main([sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32.path, "-vv", "-f", "sc32"]) == 0
-    assert capa.main.main([sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32.path, "-v", "-f", "sc32"]) == 0
-    assert capa.main.main([sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32.path, "-j", "-f", "sc32"]) == 0
-    assert capa.main.main([sample_499c2a85f6e8142c3f48d4251c9c7cd6_raw32.path, "-f", "sc32"]) == 0
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_main_shellcode(z499c2_extractor):
+    path = z499c2_extractor.path
+    assert capa.main.main([path, "-vv", "-f", "sc32"]) == 0
+    assert capa.main.main([path, "-v", "-f", "sc32"]) == 0
+    assert capa.main.main([path, "-j", "-f", "sc32"]) == 0
+    assert capa.main.main([path, "-f", "sc32"]) == 0
 
 
 def test_ruleset():
@@ -66,7 +89,7 @@ def test_ruleset():
                             name: function rule
                             scope: function
                         features:
-                          - characteristic: switch
+                          - characteristic: tight loop
                     """
                 )
             ),
@@ -89,7 +112,8 @@ def test_ruleset():
     assert len(rules.basic_block_rules) == 1
 
 
-def test_match_across_scopes_file_function(sample_9324d1a8ae37a36ae560c37448c9705a):
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_match_across_scopes_file_function(z9324d_extractor):
     rules = capa.rules.RuleSet(
         [
             # this rule should match on a function (0x4073F0)
@@ -146,16 +170,14 @@ def test_match_across_scopes_file_function(sample_9324d1a8ae37a36ae560c37448c970
             ),
         ]
     )
-    extractor = capa.features.extractors.viv.VivisectFeatureExtractor(
-        sample_9324d1a8ae37a36ae560c37448c9705a.vw, sample_9324d1a8ae37a36ae560c37448c9705a.path,
-    )
-    capabilities, meta = capa.main.find_capabilities(rules, extractor)
+    capabilities, meta = capa.main.find_capabilities(rules, z9324d_extractor)
     assert "install service" in capabilities
     assert ".text section" in capabilities
     assert ".text section and install service" in capabilities
 
 
-def test_match_across_scopes(sample_9324d1a8ae37a36ae560c37448c9705a):
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_match_across_scopes(z9324d_extractor):
     rules = capa.rules.RuleSet(
         [
             # this rule should match on a basic block (including at least 0x403685)
@@ -211,16 +233,14 @@ def test_match_across_scopes(sample_9324d1a8ae37a36ae560c37448c9705a):
             ),
         ]
     )
-    extractor = capa.features.extractors.viv.VivisectFeatureExtractor(
-        sample_9324d1a8ae37a36ae560c37448c9705a.vw, sample_9324d1a8ae37a36ae560c37448c9705a.path
-    )
-    capabilities, meta = capa.main.find_capabilities(rules, extractor)
+    capabilities, meta = capa.main.find_capabilities(rules, z9324d_extractor)
     assert "tight loop" in capabilities
     assert "kill thread loop" in capabilities
     assert "kill thread program" in capabilities
 
 
-def test_subscope_bb_rules(sample_9324d1a8ae37a36ae560c37448c9705a):
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_subscope_bb_rules(z9324d_extractor):
     rules = capa.rules.RuleSet(
         [
             capa.rules.Rule.from_yaml(
@@ -240,14 +260,12 @@ def test_subscope_bb_rules(sample_9324d1a8ae37a36ae560c37448c9705a):
         ]
     )
     # tight loop at 0x403685
-    extractor = capa.features.extractors.viv.VivisectFeatureExtractor(
-        sample_9324d1a8ae37a36ae560c37448c9705a.vw, sample_9324d1a8ae37a36ae560c37448c9705a.path,
-    )
-    capabilities, meta = capa.main.find_capabilities(rules, extractor)
+    capabilities, meta = capa.main.find_capabilities(rules, z9324d_extractor)
     assert "test rule" in capabilities
 
 
-def test_byte_matching(sample_9324d1a8ae37a36ae560c37448c9705a):
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_byte_matching(z9324d_extractor):
     rules = capa.rules.RuleSet(
         [
             capa.rules.Rule.from_yaml(
@@ -265,15 +283,12 @@ def test_byte_matching(sample_9324d1a8ae37a36ae560c37448c9705a):
             )
         ]
     )
-
-    extractor = capa.features.extractors.viv.VivisectFeatureExtractor(
-        sample_9324d1a8ae37a36ae560c37448c9705a.vw, sample_9324d1a8ae37a36ae560c37448c9705a.path,
-    )
-    capabilities, meta = capa.main.find_capabilities(rules, extractor)
+    capabilities, meta = capa.main.find_capabilities(rules, z9324d_extractor)
     assert "byte match test" in capabilities
 
 
-def test_count_bb(sample_9324d1a8ae37a36ae560c37448c9705a):
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_count_bb(z9324d_extractor):
     rules = capa.rules.RuleSet(
         [
             capa.rules.Rule.from_yaml(
@@ -292,9 +307,42 @@ def test_count_bb(sample_9324d1a8ae37a36ae560c37448c9705a):
             )
         ]
     )
-
-    extractor = capa.features.extractors.viv.VivisectFeatureExtractor(
-        sample_9324d1a8ae37a36ae560c37448c9705a.vw, sample_9324d1a8ae37a36ae560c37448c9705a.path,
-    )
-    capabilities, meta = capa.main.find_capabilities(rules, extractor)
+    capabilities, meta = capa.main.find_capabilities(rules, z9324d_extractor)
     assert "count bb" in capabilities
+
+
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_fix262(pma16_01_extractor, capsys):
+    # tests rules can be loaded successfully and all output modes
+    path = pma16_01_extractor.path
+    assert capa.main.main([path, "-vv", "-t", "send HTTP request", "-q"]) == 0
+
+    std = capsys.readouterr()
+    assert "HTTP/1.0" in std.out
+    assert "www.practicalmalwareanalysis.com" not in std.out
+
+
+@pytest.mark.xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2")
+def test_not_render_rules_also_matched(z9324d_extractor, capsys):
+    # rules that are also matched by other rules should not get rendered by default.
+    # this cuts down on the amount of output while giving approx the same detail.
+    # see #224
+    path = z9324d_extractor.path
+
+    # `act as TCP client` matches on
+    # `connect TCP client` matches on
+    # `create TCP socket`
+    #
+    # so only `act as TCP client` should be displayed
+    assert capa.main.main([path]) == 0
+    std = capsys.readouterr()
+    assert "act as TCP client" in std.out
+    assert "connect TCP socket" not in std.out
+    assert "create TCP socket" not in std.out
+
+    # this strategy only applies to the default renderer, not any verbose renderer
+    assert capa.main.main([path, "-v"]) == 0
+    std = capsys.readouterr()
+    assert "act as TCP client" in std.out
+    assert "connect TCP socket" in std.out
+    assert "create TCP socket" in std.out

tests/test_rules.py

@@ -1,11 +1,17 @@
 # Copyright (C) 2020 FireEye, 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 textwrap
 
 import pytest
 
 import capa.rules
-from capa.features import String
+from capa.features import ARCH_X32, ARCH_X64, String
 from capa.features.insn import Number, Offset
 
 
@@ -77,6 +83,16 @@ def test_rule_yaml_descriptions():
                     - string: '/SELECT.*FROM.*WHERE/i'
                       description: SQL WHERE Clause
                     - count(number(2 = AF_INET/SOCK_DGRAM)): 2
+                    - or:
+                        - and:
+                            - offset: 0x50 = IMAGE_NT_HEADERS.OptionalHeader.SizeOfImage
+                            - offset: 0x34 = IMAGE_NT_HEADERS.OptionalHeader.ImageBase
+                          description: 32-bits
+                        - and:
+                            - offset: 0x50 = IMAGE_NT_HEADERS64.OptionalHeader.SizeOfImage
+                            - offset: 0x30 = IMAGE_NT_HEADERS64.OptionalHeader.ImageBase
+                          description: 64-bits
+                      description: PE headers offsets
         """
     )
     r = capa.rules.Rule.from_yaml(rule)
@@ -87,6 +103,8 @@ def test_rule_yaml_descriptions():
                 Number(2): {2, 3},
                 String("This program cannot be run in DOS mode."): {4},
                 String("SELECT password FROM hidden_table WHERE user == admin"): {5},
+                Offset(0x50): {6},
+                Offset(0x30): {7},
             }
         )
         == True
@@ -144,6 +162,23 @@ def test_rule_yaml_count_range():
     assert r.evaluate({Number(100): {1, 2, 3}}) == False
 
 
+def test_rule_yaml_count_string():
+    rule = textwrap.dedent(
+        """
+        rule:
+            meta:
+                name: test rule
+            features:
+                - count(string(foo)): 2
+        """
+    )
+    r = capa.rules.Rule.from_yaml(rule)
+    assert r.evaluate({String("foo"): {}}) == False
+    assert r.evaluate({String("foo"): {1}}) == False
+    assert r.evaluate({String("foo"): {1, 2}}) == True
+    assert r.evaluate({String("foo"): {1, 2, 3}}) == False
+
+
 def test_invalid_rule_feature():
     with pytest.raises(capa.rules.InvalidRule):
         capa.rules.Rule.from_yaml(
@@ -249,7 +284,7 @@ def test_subscope_rules():
                                 - function:
                                     - and:
                                         - characteristic: nzxor
-                                        - characteristic: switch
+                                        - characteristic: loop
                     """
                 )
             )
@@ -362,10 +397,10 @@ def test_number_symbol():
     children = list(r.statement.get_children())
     assert (Number(1) in children) == True
     assert (Number(0xFFFFFFFF) in children) == True
-    assert (Number(2, "symbol name") in children) == True
-    assert (Number(3, "symbol name") in children) == True
-    assert (Number(4, "symbol name = another name") in children) == True
-    assert (Number(0x100, "symbol name") in children) == True
+    assert (Number(2, description="symbol name") in children) == True
+    assert (Number(3, description="symbol name") in children) == True
+    assert (Number(4, description="symbol name = another name") in children) == True
+    assert (Number(0x100, description="symbol name") in children) == True
 
 
 def test_count_number_symbol():
@@ -385,8 +420,8 @@ def test_count_number_symbol():
     assert r.evaluate({Number(2): {}}) == False
     assert r.evaluate({Number(2): {1}}) == True
     assert r.evaluate({Number(2): {1, 2}}) == False
-    assert r.evaluate({Number(0x100, "symbol name"): {1}}) == False
-    assert r.evaluate({Number(0x100, "symbol name"): {1, 2, 3}}) == True
+    assert r.evaluate({Number(0x100, description="symbol name"): {1}}) == False
+    assert r.evaluate({Number(0x100, description="symbol name"): {1, 2, 3}}) == True
 
 
 def test_invalid_number():
@@ -430,6 +465,39 @@ def test_invalid_number():
         )
 
 
+def test_number_arch():
+    r = capa.rules.Rule.from_yaml(
+        textwrap.dedent(
+            """
+            rule:
+                meta:
+                    name: test rule
+                features:
+                    - number/x32: 2
+            """
+        )
+    )
+    assert r.evaluate({Number(2, arch=ARCH_X32): {1}}) == True
+
+    assert r.evaluate({Number(2): {1}}) == False
+    assert r.evaluate({Number(2, arch=ARCH_X64): {1}}) == False
+
+
+def test_number_arch_symbol():
+    r = capa.rules.Rule.from_yaml(
+        textwrap.dedent(
+            """
+            rule:
+                meta:
+                    name: test rule
+                features:
+                    - number/x32: 2 = some constant
+            """
+        )
+    )
+    assert r.evaluate({Number(2, arch=ARCH_X32, description="some constant"): {1}}) == True
+
+
 def test_offset_symbol():
     rule = textwrap.dedent(
         """
@@ -448,10 +516,10 @@ def test_offset_symbol():
     r = capa.rules.Rule.from_yaml(rule)
     children = list(r.statement.get_children())
     assert (Offset(1) in children) == True
-    assert (Offset(2, "symbol name") in children) == True
-    assert (Offset(3, "symbol name") in children) == True
-    assert (Offset(4, "symbol name = another name") in children) == True
-    assert (Offset(0x100, "symbol name") in children) == True
+    assert (Offset(2, description="symbol name") in children) == True
+    assert (Offset(3, description="symbol name") in children) == True
+    assert (Offset(4, description="symbol name = another name") in children) == True
+    assert (Offset(0x100, description="symbol name") in children) == True
 
 
 def test_count_offset_symbol():
@@ -471,8 +539,82 @@ def test_count_offset_symbol():
     assert r.evaluate({Offset(2): {}}) == False
     assert r.evaluate({Offset(2): {1}}) == True
     assert r.evaluate({Offset(2): {1, 2}}) == False
-    assert r.evaluate({Offset(0x100, "symbol name"): {1}}) == False
-    assert r.evaluate({Offset(0x100, "symbol name"): {1, 2, 3}}) == True
+    assert r.evaluate({Offset(0x100, description="symbol name"): {1}}) == False
+    assert r.evaluate({Offset(0x100, description="symbol name"): {1, 2, 3}}) == True
+
+
+def test_offset_arch():
+    r = capa.rules.Rule.from_yaml(
+        textwrap.dedent(
+            """
+            rule:
+                meta:
+                    name: test rule
+                features:
+                    - offset/x32: 2
+            """
+        )
+    )
+    assert r.evaluate({Offset(2, arch=ARCH_X32): {1}}) == True
+
+    assert r.evaluate({Offset(2): {1}}) == False
+    assert r.evaluate({Offset(2, arch=ARCH_X64): {1}}) == False
+
+
+def test_offset_arch_symbol():
+    r = capa.rules.Rule.from_yaml(
+        textwrap.dedent(
+            """
+            rule:
+                meta:
+                    name: test rule
+                features:
+                    - offset/x32: 2 = some constant
+            """
+        )
+    )
+    assert r.evaluate({Offset(2, arch=ARCH_X32, description="some constant"): {1}}) == True
+
+
+def test_invalid_offset():
+    with pytest.raises(capa.rules.InvalidRule):
+        r = capa.rules.Rule.from_yaml(
+            textwrap.dedent(
+                """
+                rule:
+                    meta:
+                        name: test rule
+                    features:
+                        - offset: "this is a string"
+                """
+            )
+        )
+
+    with pytest.raises(capa.rules.InvalidRule):
+        r = capa.rules.Rule.from_yaml(
+            textwrap.dedent(
+                """
+                rule:
+                    meta:
+                        name: test rule
+                    features:
+                        - offset: 2=
+                """
+            )
+        )
+
+    with pytest.raises(capa.rules.InvalidRule):
+        r = capa.rules.Rule.from_yaml(
+            textwrap.dedent(
+                """
+                rule:
+                    meta:
+                        name: test rule
+                    features:
+                        - offset: symbol name = 2
+                """
+            )
+        )
 
 
 def test_invalid_string_values_int():
@@ -538,57 +680,20 @@ def test_regex_values_always_string():
         ),
     ]
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(rules), {capa.features.String("123"): {1}}, 0x0,
+        capa.engine.topologically_order_rules(rules),
+        {capa.features.String("123"): {1}},
+        0x0,
     )
     assert capa.features.MatchedRule("test rule") in features
 
     features, matches = capa.engine.match(
-        capa.engine.topologically_order_rules(rules), {capa.features.String("0x123"): {1}}, 0x0,
+        capa.engine.topologically_order_rules(rules),
+        {capa.features.String("0x123"): {1}},
+        0x0,
     )
     assert capa.features.MatchedRule("test rule") in features
 
 
-def test_invalid_offset():
-    with pytest.raises(capa.rules.InvalidRule):
-        r = capa.rules.Rule.from_yaml(
-            textwrap.dedent(
-                """
-                rule:
-                    meta:
-                        name: test rule
-                    features:
-                        - offset: "this is a string"
-                """
-            )
-        )
-
-    with pytest.raises(capa.rules.InvalidRule):
-        r = capa.rules.Rule.from_yaml(
-            textwrap.dedent(
-                """
-                rule:
-                    meta:
-                        name: test rule
-                    features:
-                        - offset: 2=
-                """
-            )
-        )
-
-    with pytest.raises(capa.rules.InvalidRule):
-        r = capa.rules.Rule.from_yaml(
-            textwrap.dedent(
-                """
-                rule:
-                    meta:
-                        name: test rule
-                    features:
-                        - offset: symbol name = 2
-                """
-            )
-        )
-
-
 def test_filter_rules():
     rules = capa.rules.RuleSet(
         [

tests/test_viv_features.py

@@ -1,306 +1,30 @@
 # Copyright (C) 2020 FireEye, 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 sys
 
-import viv_utils
-
-import capa.features
-import capa.features.file
-import capa.features.insn
-import capa.features.basicblock
-import capa.features.extractors.viv.file
-import capa.features.extractors.viv.insn
-import capa.features.extractors.viv.function
-import capa.features.extractors.viv.basicblock
 from fixtures import *
 
 
-def extract_file_features(vw, path):
-    features = set([])
-    for feature, va in capa.features.extractors.viv.file.extract_features(vw, path):
-        features.add(feature)
-    return features
-
-
-def extract_function_features(f):
-    features = collections.defaultdict(set)
-    for bb in f.basic_blocks:
-        for insn in bb.instructions:
-            for feature, va in capa.features.extractors.viv.insn.extract_features(f, bb, insn):
-                features[feature].add(va)
-        for feature, va in capa.features.extractors.viv.basicblock.extract_features(f, bb):
-            features[feature].add(va)
-    for feature, va in capa.features.extractors.viv.function.extract_features(f):
-        features[feature].add(va)
-    return features
-
-
-def extract_basic_block_features(f, bb):
-    features = set({})
-    for insn in bb.instructions:
-        for feature, _ in capa.features.extractors.viv.insn.extract_features(f, bb, insn):
-            features.add(feature)
-    for feature, _ in capa.features.extractors.viv.basicblock.extract_features(f, bb):
-        features.add(feature)
-    return features
-
-
-def test_api_features(mimikatz):
-    features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x403BAC))
-    assert capa.features.insn.API("advapi32.CryptAcquireContextW") in features
-    assert capa.features.insn.API("advapi32.CryptAcquireContext") in features
-    assert capa.features.insn.API("advapi32.CryptGenKey") in features
-    assert capa.features.insn.API("advapi32.CryptImportKey") in features
-    assert capa.features.insn.API("advapi32.CryptDestroyKey") in features
-    assert capa.features.insn.API("CryptAcquireContextW") in features
-    assert capa.features.insn.API("CryptAcquireContext") in features
-    assert capa.features.insn.API("CryptGenKey") in features
-    assert capa.features.insn.API("CryptImportKey") in features
-    assert capa.features.insn.API("CryptDestroyKey") in features
-
-
-def test_api_features_64_bit(sample_a198216798ca38f280dc413f8c57f2c2):
-    features = extract_function_features(viv_utils.Function(sample_a198216798ca38f280dc413f8c57f2c2.vw, 0x4011B0))
-    assert capa.features.insn.API("kernel32.GetStringTypeA") in features
-    assert capa.features.insn.API("kernel32.GetStringTypeW") not in features
-    assert capa.features.insn.API("kernel32.GetStringType") in features
-    assert capa.features.insn.API("GetStringTypeA") in features
-    assert capa.features.insn.API("GetStringType") in features
-    # call via thunk in IDA Pro
-    features = extract_function_features(viv_utils.Function(sample_a198216798ca38f280dc413f8c57f2c2.vw, 0x401CB0))
-    assert capa.features.insn.API("msvcrt.vfprintf") in features
-    assert capa.features.insn.API("vfprintf") in features
-
-
-def test_string_features(mimikatz):
-    features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D))
-    assert capa.features.String("SCardControl") in features
-    assert capa.features.String("SCardTransmit") in features
-    assert capa.features.String("ACR  > ") in features
-    # other strings not in this function
-    assert capa.features.String("bcrypt.dll") not in features
-
-
-def test_byte_features(sample_9324d1a8ae37a36ae560c37448c9705a):
-    features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60))
-    wanted = capa.features.Bytes(b"\xED\x24\x9E\xF4\x52\xA9\x07\x47\x55\x8E\xE1\xAB\x30\x8E\x23\x61")
-    # use `==` rather than `is` because the result is not `True` but a truthy value.
-    assert wanted.evaluate(features) == True
-
-
-def test_byte_features64(sample_lab21_01):
-    features = extract_function_features(viv_utils.Function(sample_lab21_01.vw, 0x1400010C0))
-    wanted = capa.features.Bytes(b"\x32\xA2\xDF\x2D\x99\x2B\x00\x00")
-    # use `==` rather than `is` because the result is not `True` but a truthy value.
-    assert wanted.evaluate(features) == True
-
-
-def test_number_features(mimikatz):
-    features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D))
-    assert capa.features.insn.Number(0xFF) in features
-    assert capa.features.insn.Number(0x3136B0) in features
-    # the following are stack adjustments
-    assert capa.features.insn.Number(0xC) not in features
-    assert capa.features.insn.Number(0x10) not in features
-
-
-def test_offset_features(mimikatz):
-    features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D))
-    assert capa.features.insn.Offset(0x0) in features
-    assert capa.features.insn.Offset(0x4) in features
-    assert capa.features.insn.Offset(0xC) in features
-    # the following are stack references
-    assert capa.features.insn.Offset(0x8) not in features
-    assert capa.features.insn.Offset(0x10) not in features
-
-
-def test_nzxor_features(mimikatz):
-    features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x410DFC))
-    assert capa.features.Characteristic("nzxor") in features  # 0x0410F0B
-
-
-def get_bb_insn(f, va):
-    """fetch the BasicBlock and Instruction instances for the given VA in the given function."""
-    for bb in f.basic_blocks:
-        for insn in bb.instructions:
-            if insn.va == va:
-                return (bb, insn)
-    raise KeyError(va)
-
-
-def test_is_security_cookie(mimikatz):
-    # not a security cookie check
-    f = viv_utils.Function(mimikatz.vw, 0x410DFC)
-    for va in [0x0410F0B]:
-        bb, insn = get_bb_insn(f, va)
-        assert capa.features.extractors.viv.insn.is_security_cookie(f, bb, insn) == False
-
-    # security cookie initial set and final check
-    f = viv_utils.Function(mimikatz.vw, 0x46C54A)
-    for va in [0x46C557, 0x46C63A]:
-        bb, insn = get_bb_insn(f, va)
-        assert capa.features.extractors.viv.insn.is_security_cookie(f, bb, insn) == True
-
-
-def test_mnemonic_features(mimikatz):
-    features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x40105D))
-    assert capa.features.insn.Mnemonic("push") in features
-    assert capa.features.insn.Mnemonic("movzx") in features
-    assert capa.features.insn.Mnemonic("xor") in features
-
-    assert capa.features.insn.Mnemonic("in") not in features
-    assert capa.features.insn.Mnemonic("out") not in features
-
-
-def test_peb_access_features(sample_a933a1a402775cfa94b6bee0963f4b46):
-    features = extract_function_features(viv_utils.Function(sample_a933a1a402775cfa94b6bee0963f4b46.vw, 0xABA6FEC))
-    assert capa.features.Characteristic("peb access") in features
-
-
-def test_file_section_name_features(mimikatz):
-    features = extract_file_features(mimikatz.vw, mimikatz.path)
-    assert capa.features.file.Section(".rsrc") in features
-    assert capa.features.file.Section(".text") in features
-    assert capa.features.file.Section(".nope") not in features
-
-
-def test_tight_loop_features(mimikatz):
-    f = viv_utils.Function(mimikatz.vw, 0x402EC4)
-    for bb in f.basic_blocks:
-        if bb.va != 0x402F8E:
-            continue
-        features = extract_basic_block_features(f, bb)
-        assert capa.features.Characteristic("tight loop") in features
-        assert capa.features.basicblock.BasicBlock() in features
-
-
-def test_tight_loop_bb_features(mimikatz):
-    f = viv_utils.Function(mimikatz.vw, 0x402EC4)
-    for bb in f.basic_blocks:
-        if bb.va != 0x402F8E:
-            continue
-        features = extract_basic_block_features(f, bb)
-        assert capa.features.Characteristic("tight loop") in features
-        assert capa.features.basicblock.BasicBlock() in features
-
-
-def test_file_export_name_features(kernel32):
-    features = extract_file_features(kernel32.vw, kernel32.path)
-    assert capa.features.file.Export("BaseThreadInitThunk") in features
-    assert capa.features.file.Export("lstrlenW") in features
-
-
-def test_file_import_name_features(mimikatz):
-    features = extract_file_features(mimikatz.vw, mimikatz.path)
-    assert capa.features.file.Import("advapi32.CryptSetHashParam") in features
-    assert capa.features.file.Import("CryptSetHashParam") in features
-    assert capa.features.file.Import("kernel32.IsWow64Process") in features
-    assert capa.features.file.Import("msvcrt.exit") in features
-    assert capa.features.file.Import("cabinet.#11") in features
-    assert capa.features.file.Import("#11") not in features
-
-
-def test_cross_section_flow_features(sample_a198216798ca38f280dc413f8c57f2c2):
-    features = extract_function_features(viv_utils.Function(sample_a198216798ca38f280dc413f8c57f2c2.vw, 0x4014D0))
-    assert capa.features.Characteristic("cross section flow") in features
-
-    # this function has calls to some imports,
-    # which should not trigger cross-section flow characteristic
-    features = extract_function_features(viv_utils.Function(sample_a198216798ca38f280dc413f8c57f2c2.vw, 0x401563))
-    assert capa.features.Characteristic("cross section flow") not in features
-
-
-def test_segment_access_features(sample_a933a1a402775cfa94b6bee0963f4b46):
-    features = extract_function_features(viv_utils.Function(sample_a933a1a402775cfa94b6bee0963f4b46.vw, 0xABA6FEC))
-    assert capa.features.Characteristic("fs access") in features
-
-
-def test_thunk_features(sample_9324d1a8ae37a36ae560c37448c9705a):
-    features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x407970))
-    assert capa.features.insn.API("kernel32.CreateToolhelp32Snapshot") in features
-    assert capa.features.insn.API("CreateToolhelp32Snapshot") in features
-
-
-def test_file_embedded_pe(pma_lab_12_04):
-    features = extract_file_features(pma_lab_12_04.vw, pma_lab_12_04.path)
-    assert capa.features.Characteristic("embedded pe") in features
-
-
-def test_stackstring_features(mimikatz):
-    features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x4556E5))
-    assert capa.features.Characteristic("stack string") in features
-
-
-def test_switch_features(mimikatz):
-    features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x409411))
-    assert capa.features.Characteristic("switch") in features
-
-    features = extract_function_features(viv_utils.Function(mimikatz.vw, 0x409393))
-    assert capa.features.Characteristic("switch") not in features
-
-
-def test_recursive_call_feature(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41):
-    features = extract_function_features(
-        viv_utils.Function(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, 0x10003100)
-    )
-    assert capa.features.Characteristic("recursive call") in features
-
-    features = extract_function_features(
-        viv_utils.Function(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, 0x10007B00)
-    )
-    assert capa.features.Characteristic("recursive call") not in features
-
-
-def test_loop_feature(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41):
-    features = extract_function_features(
-        viv_utils.Function(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, 0x10003D30)
-    )
-    assert capa.features.Characteristic("loop") in features
-
-    features = extract_function_features(
-        viv_utils.Function(sample_39c05b15e9834ac93f206bc114d0a00c357c888db567ba8f5345da0529cbed41.vw, 0x10007250)
-    )
-    assert capa.features.Characteristic("loop") not in features
-
-
-def test_file_string_features(sample_bfb9b5391a13d0afd787e87ab90f14f5):
-    features = extract_file_features(
-        sample_bfb9b5391a13d0afd787e87ab90f14f5.vw, sample_bfb9b5391a13d0afd787e87ab90f14f5.path,
-    )
-    assert capa.features.String("WarStop") in features  # ASCII, offset 0x40EC
-    assert capa.features.String("cimage/png") in features  # UTF-16 LE, offset 0x350E
-
-
-def test_function_calls_to(sample_9324d1a8ae37a36ae560c37448c9705a):
-    features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60))
-    assert capa.features.Characteristic("calls to") in features
-    assert len(features[capa.features.Characteristic("calls to")]) == 1
-
-
-def test_function_calls_to64(sample_lab21_01):
-    features = extract_function_features(viv_utils.Function(sample_lab21_01.vw, 0x1400052D0))  # memcpy
-    assert capa.features.Characteristic("calls to") in features
-    assert len(features[capa.features.Characteristic("calls to")]) == 8
-
-
-def test_function_calls_from(sample_9324d1a8ae37a36ae560c37448c9705a):
-    features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60))
-    assert capa.features.Characteristic("calls from") in features
-    assert len(features[capa.features.Characteristic("calls from")]) == 23
-
-
-def test_basic_block_count(sample_9324d1a8ae37a36ae560c37448c9705a):
-    features = extract_function_features(viv_utils.Function(sample_9324d1a8ae37a36ae560c37448c9705a.vw, 0x406F60))
-    assert len(features[capa.features.basicblock.BasicBlock()]) == 26
-
-
-def test_indirect_call_features(sample_a933a1a402775cfa94b6bee0963f4b46):
-    features = extract_function_features(viv_utils.Function(sample_a933a1a402775cfa94b6bee0963f4b46.vw, 0xABA68A0))
-    assert capa.features.Characteristic("indirect call") in features
-    assert len(features[capa.features.Characteristic("indirect call")]) == 3
-
-
-def test_indirect_calls_resolved(sample_c91887d861d9bd4a5872249b641bc9f9):
-    features = extract_function_features(viv_utils.Function(sample_c91887d861d9bd4a5872249b641bc9f9.vw, 0x401A77))
-    assert capa.features.insn.API("kernel32.CreatePipe") in features
-    assert capa.features.insn.API("kernel32.SetHandleInformation") in features
-    assert capa.features.insn.API("kernel32.CloseHandle") in features
-    assert capa.features.insn.API("kernel32.WriteFile") in features
+@parametrize(
+    "sample,scope,feature,expected",
+    FEATURE_PRESENCE_TESTS,
+    indirect=["sample", "scope"],
+)
+def test_viv_features(sample, scope, feature, expected):
+    with xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2"):
+        do_test_feature_presence(get_viv_extractor, sample, scope, feature, expected)
+
+
+@parametrize(
+    "sample,scope,feature,expected",
+    FEATURE_COUNT_TESTS,
+    indirect=["sample", "scope"],
+)
+def test_viv_feature_counts(sample, scope, feature, expected):
+    with xfail(sys.version_info >= (3, 0), reason="vivsect only works on py2"):
+        do_test_feature_count(get_viv_extractor, sample, scope, feature, expected)