hexpatch

#! /usr/bin/env python

#
# Copyright (c) 2015 Google Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from this
# software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from __future__ import print_function
import sys
import os
import argparse
import errno
from util import warning, error, print_to_error


# Patching operators
OP_UNKNOWN = 0
OP_AND = 1
OP_OR = 2
OP_XOR = 3
OP_REPLACE = 4
OP_VERIFY = 5
OP_COPY = 6
OP_SET = 7
operator_names = {
    "and": OP_AND,
    "or": OP_OR,
    "xor": OP_XOR,
    "replace": OP_REPLACE,
    "rep": OP_REPLACE,
    "r": OP_REPLACE,
    "verify": OP_VERIFY,
    "copy": OP_COPY,
    "set": OP_SET,
}


PATCH_HELP = \
    """Patch <offset> [and | or | xor | rep(lace)] <byte>... {verify <byte>...}
The offset from the start of the file
It may be specified as any of:
    num
    num+num
    symbol
    symbol+num
(All numbers are in hex)"""


def auto_int(x):
    # Workaround to allow hex numbers to be entered for numeric arguments
    return int(x, 16)


def validate_args(args):
    # Sanity-check the command line args and return a "valid" flag
    if not args.file:
        error("Missing the file to alter")
        return False

    if not args.map:
        warning("No map file specified")

    return True


def find_symbolic_offset(symbol_name, map_file=None):
    """ Look up a symbol in the symbol table

    Returns an offset if found, otherwise raises an exception
    """
    if not map_file:
        raise ValueError("Missing map file")

    with open(map_file, 'r') as map:
        for line in map:
            parts = line.split()
            if (parts[0] == symbol_name) and \
               (len(parts) > 1):
                try:
                    return int(parts[1], 16)
                except:
                    raise ValueError("Invalid offset:", parts[1])
        # Fell through without finding the symbol
        raise ValueError("symbol '{0:s}' not found in {1:s}".
                         format(symbol_name, map_file))


def parse_offset(offset_string, map_file=None):
    """ Parse the offset parameter into a numeric offset

    The offset may be specified as any of:
        num
        num+num
        symbol
        symbol+num

    Returns the offset, otherwise raises an exception
    """
    base_offset = 0
    offset = 0
    parts = offset_string.split("+")
    # If the first/only component isn't a hex number, view it as a symbol name
    try:
        base_offset = int(parts[0], 16)
    except:
        base_offset = find_symbolic_offset(parts[0], map_file)
    if base_offset < 0:
        raise ValueError("(Base) offset is negative")

    # Is there a 2nd part (i.e., base+offset)?
    if (len(parts) > 1):
        try:
            offset = int(parts[1], 16)
        except:
            raise ValueError("Invalid offset from base:", parts[1])

    if (base_offset + offset) < 0:
        raise ValueError("Offset is negative")
    return base_offset + offset


def parse_byte(byte_str):
    try:
        byte = int(byte_str, 16)
    except ValueError:
        print_to_error("Byte {0:s} is not a hex number".format(byte_str))
        raise
    if (byte < 0) or (byte > 0xff):
        raise ValueError("Byte {0:s} out of range".format(byte_str))
    return byte


def patch(args):
    """ Apply the args to patch the file

    Returns False if it failed (optional) verification, True if it succeeded,
    otherwise throws an exception.
    """
    try:
        size = os.path.getsize(args.file)
    except:
        raise IOError("Can't get size of patch file")

    with open(args.file, 'rb') as patch_file:
        blob = bytearray(size)
        patch_file.readinto(blob)

    # Apply each patch
    for patch in args.patch:
        if patch[0] in operator_names:
            operator = operator_names[patch[0]]
        else:
            raise ValueError("Unknown bitwise operator '{0:s}'".
                             format(patch[0]))
        base_offset = parse_offset(patch[1], args.map)

        bytes = patch[2:]
        if (operator == OP_AND) or (operator == OP_OR) or \
           (operator == OP_XOR) or (operator == OP_REPLACE) or \
           (operator == OP_VERIFY):
            # Normal operations are <op> <offset> <byte>...
            for offset, byte_str in enumerate(bytes):
                 # Normal patch processing of bytes
                byte = parse_byte(byte_str)
                if operator == OP_XOR:
                    blob[base_offset + offset] ^= byte
                elif operator == OP_OR:
                    blob[base_offset + offset] |= byte
                elif operator == OP_AND:
                    blob[base_offset + offset] &= byte
                elif operator == OP_REPLACE:
                    blob[base_offset + offset] = byte
                else:  # operator == OP_VERIFY:
                    if blob[base_offset + offset] == byte:
                        error("Verification failed")
                        return False
        elif operator == OP_COPY:
            # (Mem)Copy operations are:
            #      <op> <dst_offset> <src_offset> <count>...
            if len(bytes) != 2:
                raise ValueError("Incorrect number of copy parameters: {0:s}".
                                 format(patch))
            src_offset = parse_offset(bytes[0], args.map)
            count = int(bytes[1])
            # Assume the regions don't overlap (TODO: fix later)
            for offset in range(0, count):
                blob[base_offset + offset] = blob[src_offset + offset]
        else:  # operator == OP_SET
            # (Mem)Set operations are <op> <dst_offset> <byte> <count>...
            if len(bytes) != 2:
                raise ValueError("Incorrect number of set parameters: {0:s}".
                                 format(patch))
            byte = parse_byte(bytes[0])
            count = int(bytes[1])
            for offset in range(0, count):
                blob[base_offset + offset] = byte

    # Write the file
    if not args.out:
        args.out = args.file
    with open(args.out, 'wb') as patch_file:
        patch_file.write(blob)
    return True


def main():
    """Patch a file"""

    parser = argparse.ArgumentParser()

    # String/file args
    parser.add_argument("--file",
                        required=True,
                        help="The input file to patch")

    parser.add_argument("--map",
                        help="The .map file which provides symbolic offsets")

    parser.add_argument("--out",
                        help="The output file (default: --file")

    parser.add_argument("--patch", "-p",
                        action="append",
                        nargs='*',
                        help=PATCH_HELP)

    args = parser.parse_args()

    # Sanity-check the arguments
    if not validate_args(args):
        error("Invalid args")
        sys.exit(errno.EINVAL)

    try:
        patch(args)
    except ValueError as e:
        print_to_error("Value Error: {0}".format(e))
    except IOError as e:
        print_to_error("I/O Error: {0}".format(e))
    except:
        error("Unknown error")
        raise


## Launch main
#
if __name__ == '__main__':
    main()