diff --git a/capa/features/extractors/viv/indirect_calls.py b/capa/features/extractors/viv/indirect_calls.py index a090355a..8b04a3bc 100644 --- a/capa/features/extractors/viv/indirect_calls.py +++ b/capa/features/extractors/viv/indirect_calls.py @@ -7,7 +7,7 @@ # See the License for the specific language governing permissions and limitations under the License. import collections -from typing import TYPE_CHECKING, List, Tuple, Optional +from typing import TYPE_CHECKING, Set, List, Deque, Tuple, Union, Optional import envi import vivisect.const @@ -48,12 +48,14 @@ def get_previous_instructions(vw: VivWorkspace, va: int) -> List[int]: # ensure that it fallsthrough to this one. loc = vw.getPrevLocation(va, adjacent=True) if loc is not None: - # from vivisect.const: - # location: (L_VA, L_SIZE, L_LTYPE, L_TINFO) - (pva, _, ptype, pinfo) = vw.getPrevLocation(va, adjacent=True) + ploc = vw.getPrevLocation(va, adjacent=True) + if ploc is not None: + # from vivisect.const: + # location: (L_VA, L_SIZE, L_LTYPE, L_TINFO) + (pva, _, ptype, pinfo) = ploc - if ptype == LOC_OP and not (pinfo & IF_NOFALL): - ret.append(pva) + if ptype == LOC_OP and not (pinfo & IF_NOFALL): + ret.append(pva) # find any code refs, e.g. jmp, to this location. # ignore any calls. @@ -72,7 +74,7 @@ class NotFoundError(Exception): pass -def find_definition(vw: VivWorkspace, va: int, reg: int) -> Tuple[int, int]: +def find_definition(vw: VivWorkspace, va: int, reg: int) -> Tuple[int, Union[int, None]]: """ scan backwards from the given address looking for assignments to the given register. if a constant, return that value. @@ -88,8 +90,8 @@ def find_definition(vw: VivWorkspace, va: int, reg: int) -> Tuple[int, int]: raises: NotFoundError: when the definition cannot be found. """ - q = collections.deque() - seen = set([]) + q = collections.deque() # type: Deque[int] + seen = set([]) # type: Set[int] q.extend(get_previous_instructions(vw, va)) while q: