# Copyright 2020 Google LLC
#
# 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.

import pytest

import capa.features.address
from capa.engine import Or, And, Not, Some, Range
from capa.features.insn import Number
from capa.features.address import (
    ThreadAddress,
    ProcessAddress,
    DynamicCallAddress,
    DNTokenOffsetAddress,
    AbsoluteVirtualAddress,
    RelativeVirtualAddress,
)

ADDR1 = capa.features.address.AbsoluteVirtualAddress(0x401001)
ADDR2 = capa.features.address.AbsoluteVirtualAddress(0x401002)
ADDR3 = capa.features.address.AbsoluteVirtualAddress(0x401003)
ADDR4 = capa.features.address.AbsoluteVirtualAddress(0x401004)


def test_no_address_equality():
    no_addr = capa.features.address.NO_ADDRESS
    addr_zero = capa.features.address.AbsoluteVirtualAddress(0)

    assert no_addr == no_addr
    assert no_addr != addr_zero
    assert addr_zero != no_addr
    assert no_addr != ADDR1


def test_no_address_hash():
    no_addr = capa.features.address.NO_ADDRESS
    addr_zero = capa.features.address.AbsoluteVirtualAddress(0)

    assert hash(no_addr) != hash(addr_zero)

    s = {no_addr, addr_zero}
    assert len(s) == 2

    d = {no_addr: "no", addr_zero: "zero"}
    assert d[no_addr] == "no"
    assert d[addr_zero] == "zero"


def test_relative_address():
    with pytest.warns(DeprecationWarning):
        _ = RelativeVirtualAddress(0)


def test_dn_token_offset_address_cross_type_eq():
    addr = DNTokenOffsetAddress(0x1000, 0x10)
    assert (addr == AbsoluteVirtualAddress(0x1010)) is False
    assert (addr == "not an address") is False
    assert (addr == None) is False  # noqa: E711
    assert (addr == DNTokenOffsetAddress(0x1000, 0x10)) is True
    assert (addr == DNTokenOffsetAddress(0x1000, 0x11)) is False


def test_dn_token_offset_address_cross_type_lt():
    addr = DNTokenOffsetAddress(0x1000, 0x10)
    assert addr.__lt__(AbsoluteVirtualAddress(0x1010)) is NotImplemented
    assert addr.__lt__("not an address") is NotImplemented
    assert (addr < DNTokenOffsetAddress(0x1000, 0x11)) is True
    assert (addr < DNTokenOffsetAddress(0x1000, 0x10)) is False


def test_number():
    assert bool(Number(1).evaluate({Number(0): {ADDR1}})) is False
    assert bool(Number(1).evaluate({Number(1): {ADDR1}})) is True
    assert bool(Number(1).evaluate({Number(2): {ADDR1, ADDR2}})) is False


def test_and():
    assert bool(And([Number(1)]).evaluate({Number(0): {ADDR1}})) is False
    assert bool(And([Number(1)]).evaluate({Number(1): {ADDR1}})) is True
    assert bool(And([Number(1), Number(2)]).evaluate({Number(0): {ADDR1}})) is False
    assert bool(And([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}})) is False
    assert bool(And([Number(1), Number(2)]).evaluate({Number(2): {ADDR1}})) is False
    assert bool(And([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}, Number(2): {ADDR2}})) is True


def test_or():
    assert bool(Or([Number(1)]).evaluate({Number(0): {ADDR1}})) is False
    assert bool(Or([Number(1)]).evaluate({Number(1): {ADDR1}})) is True
    assert bool(Or([Number(1), Number(2)]).evaluate({Number(0): {ADDR1}})) is False
    assert bool(Or([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}})) is True
    assert bool(Or([Number(1), Number(2)]).evaluate({Number(2): {ADDR1}})) is True
    assert bool(Or([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}, Number(2): {ADDR2}})) is True


def test_not():
    assert bool(Not(Number(1)).evaluate({Number(0): {ADDR1}})) is True
    assert bool(Not(Number(1)).evaluate({Number(1): {ADDR1}})) is False


def test_some():
    assert bool(Some(0, [Number(1)]).evaluate({Number(0): {ADDR1}})) is True
    assert bool(Some(1, [Number(1)]).evaluate({Number(0): {ADDR1}})) is False

    assert bool(Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {ADDR1}})) is False
    assert bool(Some(2, [Number(1), Number(2), Number(3)]).evaluate({Number(0): {ADDR1}, Number(1): {ADDR1}})) is False
    assert (
        bool(
            Some(2, [Number(1), Number(2), Number(3)]).evaluate({
                Number(0): {ADDR1},
                Number(1): {ADDR1},
                Number(2): {ADDR1},
            })
        )
        is True
    )
    assert (
        bool(
            Some(2, [Number(1), Number(2), Number(3)]).evaluate({
                Number(0): {ADDR1},
                Number(1): {ADDR1},
                Number(2): {ADDR1},
                Number(3): {ADDR1},
            })
        )
        is True
    )
    assert (
        bool(
            Some(2, [Number(1), Number(2), Number(3)]).evaluate({
                Number(0): {ADDR1},
                Number(1): {ADDR1},
                Number(2): {ADDR1},
                Number(3): {ADDR1},
                Number(4): {ADDR1},
            })
        )
        is True
    )


def test_complex():
    assert True is bool(
        Or([
            And([Number(1), Number(2)]),
            Or([Number(3), Some(2, [Number(4), Number(5), Number(6)])]),
        ]).evaluate({
            Number(5): {ADDR1},
            Number(6): {ADDR1},
            Number(7): {ADDR1},
            Number(8): {ADDR1},
        })
    )

    assert False is bool(
        Or([
            And([Number(1), Number(2)]),
            Or([Number(3), Some(2, [Number(4), Number(5)])]),
        ]).evaluate({
            Number(5): {ADDR1},
            Number(6): {ADDR1},
            Number(7): {ADDR1},
            Number(8): {ADDR1},
        })
    )


def test_range():
    # unbounded range, but no matching feature
    # since the lower bound is zero, and there are zero matches, ok
    assert bool(Range(Number(1)).evaluate({Number(2): {}})) is True  # type: ignore

    # unbounded range with matching feature should always match
    assert bool(Range(Number(1)).evaluate({Number(1): {}})) is True  # type: ignore
    assert bool(Range(Number(1)).evaluate({Number(1): {ADDR1}})) is True

    # unbounded max
    assert bool(Range(Number(1), min=1).evaluate({Number(1): {ADDR1}})) is True
    assert bool(Range(Number(1), min=2).evaluate({Number(1): {ADDR1}})) is False
    assert bool(Range(Number(1), min=2).evaluate({Number(1): {ADDR1, ADDR2}})) is True

    # unbounded min
    assert bool(Range(Number(1), max=0).evaluate({Number(1): {ADDR1}})) is False
    assert bool(Range(Number(1), max=1).evaluate({Number(1): {ADDR1}})) is True
    assert bool(Range(Number(1), max=2).evaluate({Number(1): {ADDR1}})) is True
    assert bool(Range(Number(1), max=2).evaluate({Number(1): {ADDR1, ADDR2}})) is True
    assert bool(Range(Number(1), max=2).evaluate({Number(1): {ADDR1, ADDR2, ADDR3}})) is False

    # we can do an exact match by setting min==max
    assert bool(Range(Number(1), min=1, max=1).evaluate({Number(1): {}})) is False  # type: ignore
    assert bool(Range(Number(1), min=1, max=1).evaluate({Number(1): {ADDR1}})) is True
    assert bool(Range(Number(1), min=1, max=1).evaluate({Number(1): {ADDR1, ADDR2}})) is False

    # bounded range
    assert bool(Range(Number(1), min=1, max=3).evaluate({Number(1): {}})) is False  # type: ignore
    assert bool(Range(Number(1), min=1, max=3).evaluate({Number(1): {ADDR1}})) is True
    assert bool(Range(Number(1), min=1, max=3).evaluate({Number(1): {ADDR1, ADDR2}})) is True
    assert bool(Range(Number(1), min=1, max=3).evaluate({Number(1): {ADDR1, ADDR2, ADDR3}})) is True
    assert bool(Range(Number(1), min=1, max=3).evaluate({Number(1): {ADDR1, ADDR2, ADDR3, ADDR4}})) is False


def test_short_circuit():
    assert bool(Or([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}})) is True

    # with short circuiting, only the children up until the first satisfied child are captured.
    assert len(Or([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}}, short_circuit=True).children) == 1
    assert len(Or([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}}, short_circuit=False).children) == 2


def test_eval_order():
    # base cases.
    assert bool(Or([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}})) is True
    assert bool(Or([Number(1), Number(2)]).evaluate({Number(2): {ADDR1}})) is True

    # with short circuiting, only the children up until the first satisfied child are captured.
    assert len(Or([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}}).children) == 1
    assert len(Or([Number(1), Number(2)]).evaluate({Number(2): {ADDR1}}).children) == 2
    assert len(Or([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}, Number(2): {ADDR1}}).children) == 1

    # and its guaranteed that children are evaluated in order.
    assert Or([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}}).children[0].statement == Number(1)
    assert Or([Number(1), Number(2)]).evaluate({Number(1): {ADDR1}}).children[0].statement != Number(2)

    assert Or([Number(1), Number(2)]).evaluate({Number(2): {ADDR1}}).children[1].statement == Number(2)
    assert Or([Number(1), Number(2)]).evaluate({Number(2): {ADDR1}}).children[1].statement != Number(1)


def test_address_cross_type_eq():
    proc = ProcessAddress(pid=1, ppid=0)
    ava = capa.features.address.AbsoluteVirtualAddress(0x401001)

    assert (proc == ava) is False
    assert (ava == proc) is False


def test_process_address_sorting():
    proc1 = ProcessAddress(pid=1, ppid=0)
    proc2 = ProcessAddress(pid=2, ppid=0)

    assert sorted([proc2, proc1]) == [proc1, proc2]


def test_process_address_cross_type_sort_raises():
    proc = ProcessAddress(pid=1, ppid=0)
    ava = capa.features.address.AbsoluteVirtualAddress(0x401001)

    with pytest.raises(TypeError):
        sorted([proc, ava])


def test_process_address_lt_returns_not_implemented_for_other_types():
    proc = ProcessAddress(pid=1, ppid=0)
    ava = capa.features.address.AbsoluteVirtualAddress(0x401001)

    assert proc.__lt__(ava) is NotImplemented


def test_thread_address_cross_type_eq():
    proc = ProcessAddress(pid=1, ppid=0)
    thread = ThreadAddress(process=proc, tid=10)
    ava = capa.features.address.AbsoluteVirtualAddress(0x401001)

    assert (thread == ava) is False
    assert (ava == thread) is False


def test_dynamic_call_address_cross_type_eq():
    proc = ProcessAddress(pid=1, ppid=0)
    thread = ThreadAddress(process=proc, tid=10)
    call = DynamicCallAddress(thread=thread, id=0)
    ava = capa.features.address.AbsoluteVirtualAddress(0x401001)

    assert (call == ava) is False
    assert (ava == call) is False