valhalla

#!/usr/bin/python3
import re
import os
import sys
import subprocess
import traceback

import xml.etree.ElementTree as ET
from collections import namedtuple

Frame = namedtuple('Frame', ['file', 'function', 'line', 'local'])
Error = namedtuple('Error', ['xml', 'unique', 'stack', 'err_type', 'err_exp'])
CouldNotCompute = namedtuple('CouldNotCompute', ['handler', 'traceback'])

from colorama import *
init()

KIND = 'kind' 
STACK = 'stack'
ERROR = 'error'
FRAME = 'frame'
WHAT = 'what'
AUXWHAT = 'auxwhat'
UNIQUE = 'unique'

VALHALLA_HEADER = f"{Style.BRIGHT+Fore.WHITE}[VALHALLA]{Style.RESET_ALL}"

# register functions

MEMCHECK_EXPLAINATIONS = []
def register_explaination(f):
    MEMCHECK_EXPLAINATIONS.append(f)
    return f

# helper model functions

def is_local_obj(o):
    """
    Takes some string o, which represents a path.
    Determines if that path contains a "local" object 
    (that is, one written by the user).
    """
    return not any([
        o is None,
        o == "",
        "../sysdeps" in o,
        o.startswith("/usr/"),
        o.startswith("/build/"),
        ".so" in o[-6:],
        "system-suppplied" in o,
    ])

def get_stack(stack_xml):
    """
    Takes some XML node, which is a stack tag in an error
    Returns a sequence of Frame tuples which represent that stack.
    """
    stack = []
    for frame in stack_xml.findall(FRAME):
        frame_local = is_local_obj(frame.find('obj').text)
        frame_function = frame.find('fn').text
        if frame_local:
            frame_file = frame.find('file').text
            frame_line = frame.find('line').text
        else:
            frame_file = "system_file"
            frame_line = frame_function
        stack.append(Frame(frame_file, frame_function, frame_line, frame_local))
    return stack

def get_relevant_func(stack):
    """
    given a stack, return the most relevant function.
    Example:
    <random_internal_func>
    > printf
    user_function
    main
    
    the relevant function here is printf.

    Example:
    <seg fault>
    > user_function
    main
    
    the most relevant function is the user's.
    """

    foriegn_func = None
    for frame in stack:
        if frame.local and foriegn_func is None:
            user_func = frame.function
            return user_func
        elif frame.local:
            return foriegn_func
        foriegn_func = frame.function


def read_relative_locations(aux_what):
    FIELDS = ['num', 'relative', 'size', 'func']
    RelativeLocation = namedtuple('RelativeLocation', FIELDS)
    alloc_re = r"Address 0x[0-9a-f]* is (?P<num>\d*) bytes (?P<relative>\w*) a block of size (?P<size>\d*) (?P<func>\w*)'d"
    match = re.match(alloc_re, aux_what)
    if match is None:
        return match
    return RelativeLocation(*[match.group(x) for x in FIELDS])


# helper view functions

def get_location(stack):
    for frame in stack:
        if frame.local:
            return f"{frame.file}({frame.function}):{frame.line}"
    else:
        return "[No local functions]"

def pretty_print_stack(stack):
    lines = []
    first_local = True
    for frame in stack:
        line = ""
        if not frame.local:
            line += f"{Style.DIM}"
        elif frame.local and first_local:
            first_local = False
            line += f"{Fore.BLUE}->{Style.RESET_ALL}"
        line += f"\tin {frame.function} at {frame.file}:{frame.line}"
        line += f"{Style.RESET_ALL}"
        lines.append(line)
    
    lines.reverse()
    return '\n'.join(lines)

# memcheck error definitions

@register_explaination
def memcheck_overlap(error_xml):
    def match(error_xml):
        return error_xml.find(KIND).text == "Overlap"

    if not match(error_xml):
        return None
    
    unique = error_xml.find(UNIQUE).text

    stack = get_stack(error_xml.find(STACK))
    func = get_relevant_func(stack)
    err_type = "Memory Copy overlaps!"
    err_exp = f"You tried to use {func} to copy memory, but your source overlapped with your destination. Try using memmove/strmove if you wanted to overlap, or check that the memory regions don't overlap"
    return Error(error_xml, unique, stack, err_type, err_exp)


@register_explaination
def memcheck_double_free(error_xml):
    def match(error_xml):
        stacks = list(error_xml.findall(STACK))
        return error_xml.find(KIND).text == "InvalidFree" and len(stacks) == 3
    
    if not match(error_xml):
        return None

    unique = error_xml.find(UNIQUE).text

    stacks = list(error_xml.findall(STACK))
    stack, free_stack, alloc_stack = map(get_stack, stacks)
    
    err_type = "Double Free"
    err_exp = f"You tried to free memory, and it failed. The memory you freed was already freed previously, here:\n{pretty_print_stack(free_stack)}\n"
    err_exp += f"You should check to make sure your memory is only freed once. \n\nNOTE: the memory you freed was declared here:\n{pretty_print_stack(alloc_stack)}\n"
    return Error(error_xml, unique, stack, err_type, err_exp)


@register_explaination
def memcheck_invalid_free(error_xml):
    def match(error_xml):
        stacks = list(error_xml.findall(STACK))
        return error_xml.find(KIND).text == "InvalidFree" and len(stacks) > 1

    if not match(error_xml):
        return None

    unique = error_xml.find(UNIQUE).text
    
    xml_stacks = list(error_xml.findall(STACK))
    stack, alloc_stack = map(get_stack, xml_stacks[:2])
    
    match = read_relative_locations(error_xml.find(AUXWHAT).text)
    if match is None:
        return None

    err_type = "Invalid Free"
    err_exp = f"You tried to free memory, and it failed because you {match.func}'d {match.size} bytes, but the address you tried to free was {match.num} {match.relative} that. You can only call free with the exact same address you got from allocation. Check you don't modify the pointer you're freeing. {Style.BRIGHT}You allocated the memory here:{Style.RESET_ALL}\n{pretty_print_stack(alloc_stack)}"
    
    if match.num == "0" and match.relative == "inside" and len(xml_stacks) == 3:
        # means it was only a free error
        return None

    return Error(error_xml, unique, stack, err_type, err_exp)

@register_explaination
def memcheck_stack_free(error_xml):
    def match(error_xml):
        stacks = list(error_xml.findall(STACK))
        return error_xml.find(KIND).text == "InvalidFree" and "stack" in error_xml.find(AUXWHAT).text

    if not match(error_xml):
        return None

    unique = error_xml.find(UNIQUE).text
    
    stack = get_stack(error_xml.find(STACK))

    aux_whats = list(error_xml.findall(AUXWHAT))
    
    err_type = "Invalid Free On Stack"
    location = aux_whats[1].text
    err_exp = f"You tried to free memory, but that memory is defined on the stack. You can only free memory on the heap (allocated by malloc/calloc). It was defined {location}."
   
    return Error(error_xml, unique, stack, err_type, err_exp)

@register_explaination
def memcheck_fishy(error_xml):
    def match(error_xml):
        return error_xml.find(KIND).text == "FishyValue"

    if not match(error_xml):
        return None
    
    unique = error_xml.find(UNIQUE).text
    
    stack = get_stack(error_xml.find(STACK))
    
    what = error_xml.find(WHAT).text.replace('\n', '')

    err_type = "Fishy Value"
    err_exp = f"You called a system function with a value that doesn't make sense ({Style.DIM+what+Style.RESET_ALL})"
    
    return Error(error_xml, unique, stack, err_type, err_exp)


@register_explaination
def memcheck_invalid_access(error_xml):
    def match(error_xml):
        return error_xml.find(KIND).text in ["InvalidRead", "InvalidWrite"]

    if not match(error_xml):
        return None

    access_types = {
        'InvalidRead': 'read',
        'InvalidWrite': 'write',
    }
    access = access_types[error_xml.find(KIND).text]

    unique = error_xml.find(UNIQUE).text
    
    xml_stacks = error_xml.findall(STACK)
    stack, created_stack = map(get_stack, xml_stacks)
    
    match = read_relative_locations(error_xml.find(AUXWHAT).text)
    if match is None:
        return None

    err_type = "Invalid Access"
    err_exp = (f"You tried to {access} memory in a place which is not allowed (you tried to access {match.num} bytes {match.relative} the {match.size} bytes {match.func}'d).\n"
            f" Ensure you alloc enough space, or make sure you never {access} outside the space. If you are mallocing a struct, make sure you aren't mallocing a pointer to that struct.\n"
            f"{Style.BRIGHT}You allocated the memory here:{Style.RESET_ALL}\n{pretty_print_stack(created_stack)}")
    
    return Error(error_xml, unique, stack, err_type, err_exp)


@register_explaination
def memcheck_definitely_lost(error_xml):
    def match(error_xml):
        return error_xml.find(KIND).text == "Leak_DefinitelyLost"

    if not match(error_xml):
        return None

    unique = error_xml.find(UNIQUE).text
    
    stack = get_stack(error_xml.find(STACK))

    
    err_type = "Directly Lost Memory"
    err_exp = ("Your program has a memory leak. In this case, the leak is direct. \nThat means that the memory is directly accessible through a variable in code."
               "This usually indicates memory at the head of a data structure, or a string or array, that has not been free'd\n" 
               "You should free all memory that you allocate, or that a function you call allocates and returns. The place the memory is allocated is shown below.")
    
    return Error(error_xml, unique, stack, err_type, err_exp)

@register_explaination
def memcheck_indirectly_lost(error_xml):
    def match(error_xml):
        return error_xml.find(KIND).text == "Leak_IndirectlyLost"

    if not match(error_xml):
        return None

    unique = error_xml.find(UNIQUE).text
    
    stack = get_stack(error_xml.find(STACK))

    
    err_type = "Indirectly Lost Memory"
    err_exp = ("Your program has a memory leak. In this case, the leak is indirect. \nThat means that it is only referenced by another allocated block."
               "This usually indicates memory that is part of a data structure (linked list, tree, etc.) that has not been free'd\n" 
               "You should free all memory that you allocate, or that a function you call allocates and returns. The place the memory is allocated is shown below.")
    
    return Error(error_xml, unique, stack, err_type, err_exp)

@register_explaination
def memcheck_syscall_param(error_xml):
    def match(error_xml):
        return error_xml.find(KIND).text == "SyscallParam"

    if not match(error_xml):
        return None

    unique = error_xml.find(UNIQUE).text
    
    stack = get_stack(error_xml.findall(STACK)[0])
    aux_stacks = map(get_stack, error_xml.findall(STACK)[1:])

    what_text = error_xml.find(WHAT).text
    func_match = re.search(r'(?P<func>[\w_]*)\((?P<param>[\w_]*)\)', what_text)

    what_func = func_match.group('func')
    what_param = func_match.group('param')
    
    err_type = "Bad Paramater to a syscall"
    err_exp = f"You used a syscall ({what_func}), and the parameter {what_param} was not initialized. You should check that you set its value.\n"
    
    if "contains" in what_text:
        # update me
        err_exp += f"The paramater was initalized on the stack, here: \n"
        err_exp += pretty_print_stack(get_stack(error_xml.find(AUXWHAT)))
    elif "points to" in what_text:
        err_exp += f"The paramater was initalized on the heap (that is, with calloc/malloc), here: \n"
    else:
        return None

    return Error(error_xml, unique, stack, err_type, err_exp)

@register_explaination
def memcheck_uninit_cond(error_xml):
    def match(error_xml):
        return error_xml.find(KIND).text == "UninitCondition"

    if not match(error_xml):
        return None

    unique = error_xml.find(UNIQUE).text
    
    stack = get_stack(error_xml.findall(STACK)[0])
    aux_stacks = list(map(get_stack, error_xml.findall(STACK)[1:]))

    what_text = error_xml.find(WHAT).text
    auxwhat_text = error_xml.find(AUXWHAT).text
    
    err_type = "Uninitialized Condition"
   
    err_exp = ""
    aux_stack_print = pretty_print_stack(aux_stacks[0])
    if "stack" in auxwhat_text:
        err_exp += f"A paramater was declared without a value on the stack, here: \n{aux_stack_print}\n"
    elif "heap" in auxwhat_text:
        err_exp += f"A paramater was declared without a value on the heap (that is, with calloc/malloc), here: \n{aux_stack_print}\n"
    else:
        return None
    
    err_exp += f"You used it in a conditional statement (for instance if or while), here:\n"

    return Error(error_xml, unique, stack, err_type, err_exp)

@register_explaination
def memcheck_uninit_value(error_xml):
    def match(error_xml):
        return error_xml.find(KIND).text == "UninitValue"

    if not match(error_xml):
        return None

    unique = error_xml.find(UNIQUE).text
    
    stack = get_stack(error_xml.findall(STACK)[0])
    aux_stacks = list(map(get_stack, error_xml.findall(STACK)[1:]))

    what_text = error_xml.find(WHAT).text
    auxwhat_text = error_xml.find(AUXWHAT).text
    
    err_type = "Uninitialized Value"
  
    size = re.search(r'of size (\d*)', what_text)
    if size is None:
        return None
    print("MATCHED")
    size = int(size.groups()[0])

    SIZE_MAPPING = {
        1: ' (probably a char)',
        4: ' (probably an int, float, or long)',
        8: ' (probably a double)',
    }
    size_text = SIZE_MAPPING.get(size, "")
    
    aux_stack_print = pretty_print_stack(aux_stacks[0])
    err_exp = ""
    if "stack" in auxwhat_text:
        err_exp += f"A paramater of size {size} bytes{size_text} was declared without a value on the stack, here: \n{aux_stack_print}\n"
    elif "heap" in auxwhat_text:
        err_exp += f"A paramater of size {size} bytes{size_text} was declared without a value on the heap (that is, with calloc/malloc), here: \n{aux_stack_print}\n"
    else:
        return None
    
    err_exp += f"You then used it in your program, without assigning it a value, here:\n"

    return Error(error_xml, unique, stack, err_type, err_exp)

# Filtering

ALL_LINES=-1

def parse_filter_notation(filter_text):
    parts = filter_text.split(':')
    file_name = parts[0]
    lines = [-1]
    if len(parts) > 2:
        raise ValueError("Can't have more than two parts.")
    if len(parts) == 2:
        file_lines = parts[1].split('-')
        start = None
        end = None
        if len(file_lines) > 2:
            raise ValueError("Can't have more than two line endpoints")
        if len(file_lines) == 2:
            start, end = file_lines
        else:
            start, end = file_lines*2

        try:
            start = int(start)
            end = int(end)
        except ValueError:
            raise ValueError("Start and End must be integers")
        
        lines = list(range(start, end+1))

    return file_name, lines


def is_filtered(error, filter_lines):
    for frame in error.stack:
        for f in filter_lines or {}:
            if f in frame.file or f in frame.function:
                if -1 in filter_lines[f]:
                    return True
                if frame.line in filter_lines[f]:
                    return True

    return False

# View

VALGRIND_ERROR_HEADER = "=={pid}== {error}\n"

def get_valgrind_error_text(error_xml, pid):
    error_text = ""
    what_xml = error_xml.find('what')
    if what_xml is None:
        what_xml = error_xml.find('xwhat').find('text')
    error_text += VALGRIND_ERROR_HEADER.format(
        pid=pid,
        error=what_xml.text
    )
    for error in error_xml.findall(AUXWHAT):
        error_text += VALGRIND_ERROR_HEADER.format(
            pid=pid,
            error=error.text
        )
    error_text += VALGRIND_ERROR_HEADER.format(
        pid=pid,
        error="[Stack traces have been extracted and printed below.]"
    )
    return error_text

    
def get_error_count(count, total):
    pctg = (count / total * 100)
    return f"{Style.BRIGHT}{count}{Style.RESET_ALL} errors like this occured ({pctg:.2f}% of all {total} errors)"

ERROR_FORMAT = f"""
{Style.BRIGHT+Fore.RED}{{error_type}}{Style.RESET_ALL} at {{location}}
{{error}}
{Style.BRIGHT+Fore.CYAN}Explaination:{Style.RESET_ALL} {{general_explaination}}
{Style.BRIGHT+Fore.CYAN}Error occured at:{Style.RESET_ALL}
{{stack_trace}}
{{error_count}}
~~~"""

VALGRIND_XML = "/tmp/.valhalla.xml"
VALGRIND_CMD = "valgrind --leak-check=full --show-leak-kinds=all --xml=yes --xml-file={tmp} --track-origins=yes {args}"


if __name__ == "__main__":
    print(f"{Style.BRIGHT}"+"="*15+"[ VALHALLA ]"+"="*15+f"{Style.RESET_ALL}")
    print(f"{Style.DIM}Valhalla was created by @tfpk in 2018.{Style.RESET_ALL}")
    tree = None
    filter_lines = None
    VALHALLA_XML_FILE_SETTING = '--valhalla-xml-file'
    VALHALLA_INCLUDE_ONLY_SETTING = '--valhalla-filter='
    
    for arg in sys.argv:
        if VALHALLA_INCLUDE_ONLY_SETTING not in arg:
            continue
        sys.argv.remove(arg)
        arg = arg.replace(VALHALLA_INCLUDE_ONLY_SETTING, '')
        filter_lines = filter_lines or {}
        try:
            f, lines = parse_filter_notation(arg)
            filter_lines[f] = set(filter_lines.get(f, []) + lines)
        except ValueError:
            print(f"{VALHALLA_HEADER} The filter {arg} was not recognised!")
            sys.exit(1)

    if VALHALLA_XML_FILE_SETTING in sys.argv:
        print(f"{VALHALLA_HEADER} Loading from existing xml file.")
        file_index = sys.argv.index(VALHALLA_XML_FILE_SETTING) + 1
        xml_debug_file = sys.argv[file_index]
        sys.argv.remove(VALHALLA_XML_FILE_SETTING)
        sys.argv.remove(xml_debug_file)
        tree = ET.parse(xml_debug_file)
    elif '-h' in sys.argv:
        print(f"""\
{VALHALLA_HEADER} Help:
Coming Soon!
                """)
        sys.exit(0)
    else:
        args = ' '.join(sys.argv[1:])
        print(f"{VALHALLA_HEADER} running valgrind {args}.")
       
        command = VALGRIND_CMD.format(tmp=VALGRIND_XML, args=args)
        try:
            command_output = str(subprocess.check_output(command, shell=True))
        except subprocess.CalledProcessError:
            print(f"{VALHALLA_HEADER} Calling valgrind failed.")
            sys.exit(1)

        tree = ET.parse(VALGRIND_XML)

    pid = tree.find("pid").text

    errors = []

    for error_xml in tree.findall(ERROR):
        errors_found = 0
        for exp in MEMCHECK_EXPLAINATIONS:
            try:
                result = exp(error_xml)
                if result is not None:
                    errors.append(result)
                    errors_found += 1
            except Exception as e:
                result = CouldNotCompute(exp.__name__, traceback.format_exc())
                errors.append(result)
        if not errors_found:
            errors.append(error_xml)
    
    err_count = {}
    num_errors = 0
    for pair in tree.find("errorcounts").findall("pair"):
        count= int(pair.find("count").text)
        err_count[pair.find("unique").text] = count
        num_errors += count
    
    for error in errors:
        if not isinstance(error, tuple):
            print(f"{VALHALLA_HEADER} This error was not recognised by valhalla. Please report it!")
            ET.dump(error)
            continue
        if hasattr(error, 'traceback'):
            # is error
            print(f"{VALHALLA_HEADER} The handler for this error ({error.handler}) had a bug. Please report it!")
            print(error.traceback)
            continue
        if is_filtered(error, filter_lines):
            continue

        format_dict = {}
        format_dict['error_type'] = error.err_type
        format_dict['location'] = get_location(error.stack)
        format_dict['error'] = get_valgrind_error_text(error.xml, pid)
        format_dict['general_explaination'] = error.err_exp
        format_dict['stack_trace'] = pretty_print_stack(error.stack)
        if error.unique in err_count:
            format_dict['error_count'] = get_error_count(err_count[error.unique], num_errors)
        else:
            format_dict['error_count'] = ""

        print(ERROR_FORMAT.format(**format_dict))