test.py

#!/usr/bin/python3

__unittest = True

import subprocess
import re
import unittest
import xmlrunner
from subprocess import CalledProcessError, Popen, PIPE

#-------helpers---------------

def starts_uint( s ):
    matches = re.findall('^\d+', s)
    if matches:
        return (int(matches[0]), len(matches[0]))
    else:
        return (0, 0)

def starts_int( s ):
    matches = re.findall('^-?\d+', s)
    if matches:
        return (int(matches[0]), len(matches[0]))
    else:
        return (0, 0)

def unsigned_reinterpret(x):
    if x < 0:
        return x + 2**64
    else:
        return x

def first_or_empty( s ):
    sp = s.split()
    if sp == [] :
        return ''
    else:
        return sp[0]

#-----------------------------

before_all="""
%macro call 1
mov rax, -1
push rbx
push rbp
push r12 
push r13 
push r14 
push r15 
call %1
cmp r15, [rsp] 
jne convention_error
pop r15
cmp r14, [rsp] 
jne convention_error
pop r14
cmp r13, [rsp] 
jne convention_error
pop r13
cmp r12, [rsp] 
jne convention_error
pop r12
cmp rbp, [rsp] 
jne convention_error
pop rbp
cmp rbx, [rsp] 
jne convention_error
pop rbx
mov rdi, -1
mov rsi, -1
mov rcx, -1
mov r8, -1
mov r9, -1
mov r10, -1
mov r11, -1
%endmacro

%include "lib.inc"

global _start

section .text
convention_error:
    mov rax, 1
    mov rdi, 2
    mov rsi, err_calling_convention
    mov rdx, err_calling_convention.end - err_calling_convention
    syscall
    mov rax, 60
    mov rdi, -41
    syscall

section .data
    err_calling_convention: db "You did not respect the calling convention! Check that you handled caller-saved and callee-saved registers correctly", 10
    .end:
"""

class IOLibraryTest(unittest.TestCase):
    def compile(self, fname, text):
        f = open( fname + '.asm', 'w')
        f.write( text )
        f.close()

        self.assertEqual(subprocess.call( ['nasm', '-f', 'elf64', fname + '.asm', '-o', fname+'.o'] ), 0, 'failed to compile')
        self.assertEqual(subprocess.call( ['ld', '-o' , fname, fname+'.o'] ), 0, 'failed to link')

    def launch(self, fname, input):
        output = b''
        try:
            p = Popen(['./'+fname], shell=None, stdin=PIPE, stdout=PIPE)
            (output, _) = p.communicate(input.encode())
            self.assertNotEqual(p.returncode, -11, 'segmentation fault')
            return (output.decode(), p.returncode)
        except CalledProcessError as exc:
            self.assertNotEqual(exc.returncode, -11, 'segmentation fault')
            return (exc.output.decode(), exc.returncode)

    def perform(self, fname, text, input):
        self.compile(fname, before_all + text)
        return self.launch(fname, input)



    def test_string_length(self):
        inputs = ['asdkbasdka', 'qwe qweqe qe', '']
        for input in inputs:
            text = """
section .data
str: db '""" + input + """', 0

section .text
_start:
    mov rdi, str
    call string_length
    mov rdi, rax
    mov rax, 60
    syscall
"""
            (output, code) = self.perform('string_length', text, input)
            self.assertEqual(code, len(input), 'string_length(%s) returned wrong length: %d' % (repr(input), code))



    def test_print_string(self):
        inputs = ['ashdb asdhabs dahb', ' ', '']
        for input in inputs:
            text = """
section .data
str: db '""" + input + """', 0

section .text
_start:
    mov rdi, str
    call print_string
    xor rdi, rdi
    mov rax, 60
    syscall
"""
            (output, code) = self.perform('print_string', text, input)
            self.assertEqual(output, input, 'print_string(%s) printed wrong string: %s' % (repr(input), repr(output)))



    def test_string_copy(self):
        inputs = ['ashdb asdhabs dahb', ' ', '']
        for input in inputs:
            text = """
section .data
arg1: db '""" + input + """', 0
arg2: times """ + str(len(input) + 1) + """ db  66

section .text
_start:
    mov rdi, arg1
    mov rsi, arg2
    mov rdx, """ + str(len(input) + 1) + """
    call string_copy

    mov rdi, arg2 
    call print_string
    mov rax, 60
    xor rdi, rdi
    syscall
"""
            (output, code) = self.perform('string_copy', text, input)
            self.assertEqual(output, input, 'string_copy(%s) put wrong string into buffer: %s' % (repr(input), repr(output)))



    def test_string_copy_too_long(self):
        inputs = ['ashdb asdhabs dahb', ' ', '']
        for input in inputs:
            text = """
section .rodata
err_too_long_msg: db "string is too long", 10, 0

section .data
arg1: db '""" + input + """', 0
arg2: times """ + str(len(input)//2) + """ db  66

section .text
_start:
    mov rdi, arg1
    mov rsi, arg2
    mov rdx, """ + str(len(input)//2) + """
    call string_copy
    test rax, rax
    jnz .good
    mov rdi, err_too_long_msg 
    call print_string
    jmp _exit
    .good:
    mov rdi, arg2 
    call print_string
_exit:
    mov rax, 60
    xor rdi, rdi
    syscall
"""
            (output, code) = self.perform('string_copy_too_long', text, input)
            self.assertNotEqual(output.find('too long'), -1, 'string_copy(%s) should have failed, but returned: %s' % (repr(input), repr(output)))



    def test_print_char(self):
        inputs = ['a', ' ', 'c']
        for input in inputs:
            text = """
section .text
_start:
    mov rdi, '""" + input + """'
    call print_char
    mov rax, 60
    xor rdi, rdi
    syscall
"""
            (output, code) = self.perform('print_char', text, input)
            self.assertEqual(output, input, 'print_char(%s) printed wrong char: %s' % (repr(input), repr(output)))



    def test_print_uint(self):
        inputs = ['-1', '12345234121', '0', '12312312', '123123']
        for input in inputs:
            text = """
section .text
_start:
    mov rdi, """ + input + """
    call print_uint
    mov rax, 60
    xor rdi, rdi
    syscall
"""
            (output, code) = self.perform('print_uint', text, input)
            uinput = str(unsigned_reinterpret(int(input)))
            self.assertEqual(output, uinput, 'print_uint(%s) printed wrong number: %s, expected: %s' % (repr(input), repr(output), repr(uinput)))



    def test_print_int(self):
        inputs = ['-1', '-12345234121', '0', '123412312', '123123']
        for input in inputs:
            text = """
section .text
_start:
    mov rdi, """ + input + """
    call print_int
    mov rax, 60
    xor rdi, rdi
    syscall
"""
            (output, code) = self.perform('print_int', text, input)
            self.assertEqual(output, input, 'print_int(%s) printed wrong number: %s' % (repr(input), repr(output)))



    def test_read_char(self):
        inputs = ['-1', '-1234asdasd5234121', '', '   ', '\t   ', 'hey ya ye ya', 'hello world' ]
        for input in inputs:
            text = """
section .text
_start:
    call read_char
    mov rdi, rax
    mov rax, 60
    syscall
"""
            (output, code) = self.perform('read_char', text, input)
            if input == "":
                self.assertEqual(code, 0, 'read_char with empty input should return 0')
            else:
                self.assertEqual(code, ord(input[0]), 'read_char(%d) returned incorrect char: %d' % (ord(input[0]), code))



    def test_read_word(self):
        inputs = ['-1'] # , '-1234asdasd5234121', '', '   ', '\t   ', 'hey ya ye ya', 'hello world' ],

        for input in inputs:
            text = """
section .data
word_buf: times 20 db 0xca

section .text
_start:
    mov rdi, word_buf
    mov rsi, 20 
    call read_word
    mov rdi, rax
    call print_string

    mov rax, 60
    xor rdi, rdi
    syscall
"""
            (output, code) = self.perform('read_word', text, input)
            input_word = first_or_empty(input)
            self.assertEqual(output, input_word, 'read_word(%s) put incorrect word in the buffer: %s, expected: %s' % (repr(input), repr(output), repr(input_word)))



    def test_read_word_length(self):
        inputs = ['-1', '-1234asdasd5234121', '', '   ', '\t   ', '\t   123', 'hey ya ye ya', 'hello world' ]
        for input in inputs:
            text = """
section .data
word_buf: times 20 db 0xca

section .text
_start:
    mov rdi, word_buf
    mov rsi, 20 
    call read_word
    mov rax, 60
    mov rdi, rdx
    syscall
"""
            (output, code) = self.perform('read_word_length', text, input)
            input_word = first_or_empty(input)
            self.assertEqual(code, len(input_word), 'read_word(%s) returned incorrect length: %d, expected: %d' % (repr(input), code, len(input_word)))



    def test_read_word_too_long(self):
        inputs = [ 'asdbaskdbaksvbaskvhbashvbasdasdads wewe', 'short' ]
        for input in inputs:
            text = """
section .data
stub: times 5 db 0xca
word_buf: times 20 db 0xca

section .text
_start:
    mov rdi, word_buf
    mov rsi, 20 
    call read_word
    mov rdi, rax
    mov rax, 60
    syscall
"""
            (output, code) = self.perform('read_word_too_long', text, input)
            input_word = first_or_empty(input)
            if len(input_word) > 19:
                self.assertEqual(code, 0, 'read_word(%s) overflows buffer, but does not fail' % repr(input))
            else:
                self.assertNotEqual(code, 0, 'read_word(%s) does not overflow buffer, but fails' % repr(input))



    def test_parse_uint(self):
        inputs = ["0", "1234567890987654321hehehey", "1" ]
        for input in inputs:
            text = """
section .data
input: db '""" + input  + """', 0

section .text
_start:
    mov rdi, input
    call parse_uint
    push rdx
    mov rdi, rax
    call print_uint
    mov rax, 60
    pop rdi
    syscall
"""
            (output, code) = self.perform('parse_uint', text, input)
            (input_num, input_len) = starts_uint(input)

            self.assertEqual(output, str(input_num), 'parse_uint(%s) parsed wrong number: %s, expected: %s' % (repr(input), repr(output), repr(str(input_num))))
            self.assertEqual(code, input_len, 'parse_uint(%s) returned wrong length: %d, expected: %d' % (repr(input), code, input_len))



    def test_parse_int(self):
        inputs = ["0", "1234567890987654321hehehey", "-1dasda", "-eedea", "-123123123", "1" ]
        for input in inputs:
            text = """
section .data
input: db '""" + input  + """', 0

section .text
_start:
    mov rdi, input
    call parse_int
    push rdx
    mov rdi, rax
    call print_int
    mov rax, 60
    pop rdi
    syscall
"""
            (output, code) = self.perform('parse_int', text, input)
            (input_num, input_len) = starts_int(input)

            if input_len == 0:
                self.assertEqual(output, '0', 'parse_int(%s) should have failed, but parsed %s' % (repr(input), output))
            else:
                self.assertEqual(output, str(input_num), 'parse_int(%s) parsed wrong number: %s, expected: %s' % (repr(input), repr(output), repr(str(input_num))))
                self.assertEqual(code, input_len, 'parse_int(%s) returned wrong length: %d, expected: %d' % (repr(input), code, input_len))



    def test_string_equals(self):
        inputs = ['ashdb asdhabs dahb', ' ', '', "asd" ]
        for input in inputs:
            text = """
section .data
str1: db '""" + input + """',0
str2: db '""" + input + """',0

section .text
_start:
    mov rdi, str1
    mov rsi, str2
    call string_equals
    mov rdi, rax
    mov rax, 60
    syscall
"""
            (output, code) = self.perform('string_equals', text, input)
            self.assertEqual(code, 1, 'string_equals(%s, %s) should return 1' % (repr(input), repr(input)))



    def test_string_not_equals(self):
        inputs = ['ashdb asdhabs dahb', ' ', '', "asd" ]
        for input in inputs:
            text = """
section .data
str1: db '""" + input + """',0
str2: db '""" + input + """!!',0

section .text
_start:
    mov rdi, str1
    mov rsi, str2
    call string_equals
    mov rdi, rax
    mov rax, 60
    syscall
"""
            (output, code) = self.perform('string_not_equals', text, input)
            self.assertEqual(code, 0, 'string_equals(%s, %s!!) should return 0' % (repr(input), repr(input)))



if __name__ == "__main__":
    with open('report.xml', 'w') as report:
        unittest.main(testRunner=xmlrunner.XMLTestRunner(output=report), failfast=False, buffer=False, catchbreak=False)