# Copyright (c) 2017 Trail of Bits, 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.

import collections
import idaapi
import idautils
import idc
import itertools
import struct
import inspect
import ida_ua
import ida_bytes

_DEBUG_FILE = None
_DEBUG_PREFIX = ""
_INFO = idaapi.get_inf_structure()

# Map of the external functions names which does not return to a tuple containing information
# like the number of agruments and calling convention of the function.
_NORETURN_EXTERNAL_FUNC = {}

FUNC_LSDA_ENTRIES = collections.defaultdict()

IS_ARM = "ARM" in _INFO.procName

# True if we are running on an ELF file.
IS_ELF = (idaapi.f_ELF == _INFO.filetype) or \
         (idc.GetLongPrm(idc.INF_FILETYPE) == idc.FT_ELF)

# True if this is a Windows PE file.
IS_PE = idaapi.f_PE == _INFO.filetype

if IS_ARM:
  from arm_util import *
else:
  from x86_util import *

def INIT_DEBUG_FILE(file):
  global _DEBUG_FILE
  _DEBUG_FILE = file

def DEBUG_PUSH():
  global _DEBUG_PREFIX
  _DEBUG_PREFIX += "  "

def DEBUG_POP():
  global _DEBUG_PREFIX
  _DEBUG_PREFIX = _DEBUG_PREFIX[:-2]

def DEBUG(s):
  global _DEBUG_FILE
  if _DEBUG_FILE:
    _DEBUG_FILE.write("{}{}\n".format(_DEBUG_PREFIX, str(s)))

# Python 2.7's xrange doesn't work with `long`s.
def xrange(begin, end=None, step=1):
  if end:
    return iter(itertools.count(begin, step).next, end)
  else:
    return iter(itertools.count().next, begin)

_NOT_INST_EAS = set()

# sign extension to the given bits
def sign_extend(x, b):
  m = 1 << (b - 1)
  x = x & ((1 << b) - 1)
  return (x ^ m) - m

# Returns `True` if `ea` belongs to some code segment.
#
# TODO(pag): This functon is extra aggressive, in that it doesn't strictly
#            trust the `idc.is_code`. I have observed cases where data in
#            `.bss` is treated as code and I am not sure why. Perhaps adding
#            a reference to the data did this.
#
#            I think it has something to do with ELF thunks, e.g. entries in
#            the `.plt` section. When I made this function stricter,
#            `mcsema-lift` would report issues where it needed to add tail-calls
#            to externals.
def is_code(ea):
  if is_invalid_ea(ea):
    return False

  seg_ea = idc.get_segm_start(ea)
  seg_type = idc.get_segm_attr(seg_ea, idc.SEGATTR_TYPE)
  return (seg_type == idc.SEG_CODE)

# A stricter form of `is_code`, where we also check whether IDA thinks something
# is code. IDA is able to identify some things like embedded exception tables
# in the code section as not truly being code.
def is_code_by_flags(ea):
  if not is_code(ea):
    return False

  flags = idc.get_full_flags(ea)
  return idc.is_code(flags)

def is_read_only_segment(ea):
  mask_perms = idaapi.SEGPERM_WRITE | idaapi.SEGPERM_READ
  perms = idc.get_segm_attr(ea, idc.SEGATTR_PERM)
  return idaapi.SEGPERM_READ == (perms & mask_perms)

def is_tls_segment(ea):
  try:
    seg_name = idc.get_segm_name(ea)
    return seg_name in (".tbss", ".tdata", ".tls")
  except:
    return False

# Returns `True` if `ea` looks like a thread-local thing.
def is_tls(ea):
  if is_invalid_ea(ea):
    return False

  if is_tls_segment(ea):
    return True

  # Something references `ea`, and that something is commented as being a
  # `TLS-reference`. This comes up if you have an thread-local extern variable
  # declared/used in a binary, and defined in a shared lib. There will be an
  # offset variable.
  for source_ea in _drefs_to(ea):
    comment = idc.GetCommentEx(source_ea, 0)
    if isinstance(comment, str) and "TLS-reference" in comment:
      return True

  return False

# Mark an address as containing code.
def try_mark_as_code(ea):
  if is_code(ea) and not is_code_by_flags(ea):
    idc.MakeCode(ea)
    idaapi.auto_wait()
    return True
  return False

def mark_as_not_code(ea):
  global _NOT_INST_EAS
  _NOT_INST_EAS.add(ea)

def read_bytes_slowly(start, end):
  bytestr = []
  for i in xrange(start, end):
    if idc.has_value(idc.get_full_flags(i)):
      bt = idc.get_wide_byte(i)
      bytestr.append(chr(bt))
    else:
      bytestr.append("\x00")
  return "".join(bytestr)

def read_byte(ea):
  byte = read_bytes_slowly(ea, ea + 1)
  byte = ord(byte) 
  return byte

def read_word(ea):
  bytestr = read_bytes_slowly(ea, ea + 2)
  word = struct.unpack("<L", bytestr)[0]
  return word

def read_dword(ea):
  bytestr = read_bytes_slowly(ea, ea + 4)
  dword = struct.unpack("<L", bytestr)[0]
  return dword

def read_qword(ea):
  bytestr = read_bytes_slowly(ea, ea + 8)
  qword = struct.unpack("<Q", bytestr)[0]
  return qword

def read_leb128(ea, signed):
  """ Read LEB128 encoded data
  """
  val = 0
  shift = 0
  while True:
    byte = idc.get_wide_byte(ea)
    val |= (byte & 0x7F) << shift
    shift += 7
    ea += 1
    if (byte & 0x80) == 0:
      break

    if shift > 64:
      DEBUG("Bad leb128 encoding at {0:x}".format(ea - shift/7))
      return idc.BADADDR

  if signed and (byte & 0x40):
    val -= (1<<shift)
  return val, ea

def read_pointer(ea):
  if _INFO.is_64bit():
    return read_qword(ea)