#!/usr/bin/python
# Copyright (c) 2013 The Native Client Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

# Executable specification of valid instructions and superinstructions (in terms
# of their disassembler listing).
# Should serve as formal and up-to-date ABI reference and as baseline for
# validator exhaustive tests.

# It is generally organized as a set of functions responsible for recognizing
# and validating specific patterns (jump instructions, regular instructions,
# superinstructions, etc.)
# There are three outcomes for running such function:
#   - function raises DoNotMatchError (which means instruction is of completely
#     different structure, for example when we call ValidateSuperinstruction on
#     nop)
#   - function raises SandboxingError (which means instruction generally matches
#     respective pattern, but some rules are violated)
#   - function returns (which means instruction(s) is(are) safe)
#
# Why exceptions instead of returning False or something? Because they carry
# stack traces, which makes it easier to investigate why particular instruction
# was rejected.
# Why distinguish DoNotMatchError and SandboxingError? Because on the topmost
# level we attempt to call all matchers and we need to see which error message
# was most relevant.

import re


class DoNotMatchError(Exception):
  pass


class SandboxingError(Exception):
  pass


BUNDLE_SIZE = 32


def _ValidateLongNop(instruction):
  # Short nops do not require special exceptions (such as allowing repeated
  # prefixes and segment access), so they are handled as regular instructions.
  if re.match(r'nopw 0x0\(%[er]ax,%[er]ax,1\)$',
      instruction.disasm):
    return
  if re.match(
      r'(data32 )*nopw %cs:0x0\(%[er]ax,%[er]ax,1\)$',
      instruction.disasm):
    return
  raise DoNotMatchError(instruction)


def _ValidateStringInstruction(instruction):
  prefix_re = r'(rep |repz |repnz )?'
  lods_re = r'lods %ds:\(%esi\),(%al|%ax|%eax)'
  stos_re = r'stos (%al|%ax|%eax),%es:\(%edi\)'
  scas_re = r'scas %es:\(%edi\),(%al|%ax|%eax)'
  movs_re = r'movs[bwl] %ds:\(%esi\),%es:\(%edi\)'
  cmps_re = r'cmps[bwl] %es:\(%edi\),%ds:\(%esi\)'

  string_insn_re = '%s(%s)$' % (
      prefix_re,
      '|'.join([lods_re, stos_re, scas_re, movs_re, cmps_re]))

  if re.match(string_insn_re, instruction.disasm):
    return

  raise DoNotMatchError(instruction)


def _ValidateTlsInstruction(instruction):
  if re.match(r'mov %gs:(0x0|0x4),%e[a-z][a-z]$', instruction.disasm):
    return

  raise DoNotMatchError(instruction)


# What can follow 'j' in conditional jumps 'je', 'jno', etc.
_CONDITION_SUFFIX_RE = r'(a(e?)|b(e?)|g(e?)|l(e?)|(n?)e|(n?)o|(n?)p|(n?)s)'


def _AnyRegisterRE(group_name='register'):
  # TODO(shcherbina): explicitly list all kinds of registers we care to
  # distinguish for validation purposes.
  return r'(?P<%s>%%(st\(\d+\)|\w+))' % group_name


def _HexRE(group_name='value'):
  return r'(?P<%s>-?0x[\da-f]+)' % group_name


def _ImmediateRE(group_name='immediate'):
  return r'(?P<%s>\$%s)' % (
      group_name,
      _HexRE(group_name=group_name + '_value'))


def MemoryRE(group_name='memory'):
  # Possible forms:
  #   (%eax)
  #   (%eax,%ebx,1)
  #   (,%ebx,1)
  #   0x42(...)
  # and even
  #   0x42
  return r'(?P<%s>(?P<%s_segment>%%[cdefgs]s:)?%s?(\(%s?(,%s,\d)?\))?)' % (
      group_name,
      group_name,
      _HexRE(group_name=group_name + '_offset'),
      _AnyRegisterRE(group_name=group_name + '_base'),
      _AnyRegisterRE(group_name=group_name + '_index'))


def _IndirectJumpTargetRE(group_name='target'):
  return r'(?P<%s>\*(%s|%s))' % (
      group_name,
      _AnyRegisterRE(group_name=group_name + '_register'),
      MemoryRE(group_name=group_name + '_memory'))


def _OperandRE(group_name='operand'):
  return r'(?P<%s>%s|%s|%s|%s)' % (
      group_name,
      _AnyRegisterRE(group_name=group_name + '_register'),
      _ImmediateRE(group_name=group_name + '_immediate'),
      MemoryRE(group_name=group_name + '_memory'),
      _IndirectJumpTargetRE(group_name=group_name + '_target'))


def _SplitOps(insn, args):
  # We can't use just args.split(',') because operands can contain commas
  # themselves, for example '(%r15,%rax,1)'.
  ops = []
  i = 0
  while True:
    # We do not use mere re.match(_OperandRE(), args, i) here because
    # python backtracking regexes do not guarantee to find longest match.
    m = re.compile(r'(%s)($|,)' % _OperandRE()).match(args, i)
    assert m is not None, (args, i)
    ops.append(m.group(1))
    i = m.end(1)
    if i == len(args):
      break
    assert args[i] == ',', (insn, args, i)
    i += 1
  return ops


def ParseInstruction(instruction):
  """Parse an instruction into operands.

  Args:
    instruction: objdump_parser.Instruction tuple
  Returns:
    prefixes, mnemonic, operands
  Raises:
    SandboxingError on erroneous bytes.
  """
  # Strip comment.
  disasm, _, _ = instruction.disasm.partition('#')
  elems = disasm.split()

  if elems == []:
    raise SandboxingError(
        'disasm is empty', instruction)

  prefixes = []
  while elems != [] and elems[0] in [
      'lock', 'rep', 'repz', 'repnz',
      'data16', 'data32', 'addr16', 'addr32', 'addr64']:
    prefixes.append(elems.pop(0))

  if elems == []:
    raise SandboxingError(
        'dangling legacy prefixes', instruction)

  name = elems[0]

  if re.match(r'rex([.]W?R?X?B?)?$', name):
    raise SandboxingError('dangling rex prefix', instruction)

  # There could be branching expectation information in instruction names:
  #    jo,pt      <addr>
  #    jge,pn     <addr>
  name_re = r'[a-z]\w*(,p[nt])?$'
  assert re.match(name_re, name) or name == "nop/reserved", name

  if len(elems) == 1:
    ops = []
  elif len(elems) == 2:
    ops = _SplitOps(instruction, elems[1])
  else:
    assert False, instruction

  return prefixes, name, ops


REG32_TO_REG64 = {
    '%eax' : '%rax',
    '%ebx' : '%rbx',
    '%ecx' : '%rcx',
    '%edx' : '%rdx',
    '%esi' : '%rsi',
    '%edi' : '%rdi',
    '%esp' : '%rsp',
    '%ebp' : '%rbp',
    '%r8d' : '%r8',
    '%r9d' : '%r9',
    '%r10d' : '%r10',
    '%r11d' : '%r11',
    '%r12d' : '%r12',
    '%r13d' : '%r13',
    '%r14d' : '%r14',
    '%r15d' : '%r15'}

REGS32 = REG32_TO_REG64.keys()
REGS64 = REG32_TO_REG64.values()


class Condition(object):
  """Represents assertion about the state of 64-bit registers.

  (used as precondition and postcondition)

  Supported assertions:
    0. %rpb and %rsp are sandboxed (and nothing is known about other registers)
    1. {%rax} is restricted, %rbp and %rsp are sandboxed
    2-13. same for %rbx-%r14 not including %rbp and %rsp
    14. %rbp is restricted, %rsp is sandboxed
    15. %rsp is restricted, %rpb is sandboxed

  It can be observed that all assertions 1..15 differ from default 0 in a single
  register, which prompts internal representation of a single field,
  _restricted_register, which stores name of this standing out register
  (or None).

  * 'restricted' means higher 32 bits are zeroes
  * 'sandboxed' means within [%r15, %r15 + 2**32) range
  It goes without saying that %r15 is never changed and by definition sandboxed.
  """

  def __init__(self, restricted=None, restricted_instead_of_sandboxed=None):
    self._restricted_register = None
    if restricted is not None:
      assert restricted_instead_of_sandboxed is None
      assert restricted in REGS64
      assert restricted not in ['%r15', '%rbp', '%rsp']
      self._restricted_register = restricted
    if restricted_instead_of_sandboxed is not None:
      assert restricted is None
      assert restricted_instead_of_sandboxed in ['%rbp', '%rsp']
      self._restricted_register = restricted_instead_of_sandboxed

  def GetAlteredRegisters(self):
    """ Return pair (restricted, restricted_instead_of_sandboxed).

    Each item is either register name or None.
    """
    if self._restricted_register is None:
      return None, None