diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py new file mode 100644 index 00000000..69a79bd6 --- /dev/null +++ b/tests/test_optimizer.py @@ -0,0 +1,65 @@ +# Copyright (C) 2021 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 +import capa.engine +import capa.optimizer +import capa.features.common +from capa.engine import Or, And +from capa.features.insn import Mnemonic +from capa.features.common import Arch, Bytes, Substring + + +def test_optimizer_order(): + rule = textwrap.dedent( + """ + rule: + meta: + name: test rule + scope: function + features: + - and: + - substring: "foo" + - arch: amd64 + - mnemonic: cmp + - and: + - bytes: 3 + - offset: 2 + - or: + - number: 1 + - offset: 4 + """ + ) + r = capa.rules.Rule.from_yaml(rule) + + # before optimization + children = list(r.statement.get_children()) + assert isinstance(children[0], Substring) + assert isinstance(children[1], Arch) + assert isinstance(children[2], Mnemonic) + assert isinstance(children[3], And) + assert isinstance(children[4], Or) + + # after optimization + capa.optimizer.optimize_rules([r]) + children = list(r.statement.get_children()) + + # cost: 0 + assert isinstance(children[0], Arch) + # cost: 1 + assert isinstance(children[1], Mnemonic) + # cost: 2 + assert isinstance(children[2], Substring) + # cost: 3 + assert isinstance(children[3], Or) + # cost: 4 + assert isinstance(children[4], And)