test.py 5.07 KB
Newer Older
Enrico Pozzobon committed
1 2
#!/usr/bin/env python3

Enrico Pozzobon committed
3
import re
Enrico Pozzobon committed
4 5 6 7 8 9 10
import os
import sys
import struct
from subprocess import Popen, PIPE


def main(argv):
Enrico Pozzobon committed
11 12 13
    if len(argv) < 3:
        print("Usage: test.py LWC_AEAD_KAT.txt program [arguments]")

14 15 16 17 18
    cmd = argv[2:]

    for attempt in range(3):
        print("beginning test %d of '%s' using test vectors '%s'" % (attempt, ' '.join(cmd), argv[1]))
        try:
19 20 21 22 23
            measurements = begin_measurement()
            try:
                test(argv[1], cmd)
            finally:
                end_measurement(measurements)
24 25
            print("TEST SUCCESSFUL")
            return 0
26

27 28 29 30 31 32 33 34 35 36 37
        except Exception as ex:
            print(str(ex))
            print("TEST FAILED")
        finally:
            sys.stdout.flush()
            sys.stderr.flush()

    return 1

def test(test_file, cmd):

Enrico Pozzobon committed
38
    test_file = open(test_file, 'r')
39
    p = Popen(cmd, bufsize=0, stdin=PIPE, stdout=PIPE)
Enrico Pozzobon committed
40 41 42 43 44 45 46

    def write(data):
        l = p.stdin.write(data)
        if len(data) != l:
            raise Exception("could not write %d bytes of data (put %d)" % (len(data), l))

    def read(l):
47 48
        if l == 0:
            return b""
Enrico Pozzobon committed
49
        data = p.stdout.read(l)
50 51
        if len(data) == 0:
            print("Unexpected end of stream", file=sys.stderr)
52
            #sys.exit(1)
Enrico Pozzobon committed
53 54 55 56 57 58 59 60 61 62 63 64 65 66
        if len(data) != l:
            raise Exception("could not read %d bytes of data (got %d)" % (l, len(data)))
        return data

    def submit(action, data):
        h = struct.pack("<BI", ord(action), len(data))
        write(h)
        write(data)

    def obtain():
        l = read(4)
        (l, ) = struct.unpack("<I", l)
        return read(l)

67 68 69
    output = read(14)
    if output != b"Hello, World!\n":
        raise Exception("Unexpected output: %s" % output)
70
    print("Ready")
Enrico Pozzobon committed
71

Enrico Pozzobon committed
72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131
    m = b""
    ad = b""
    k = b""
    npub = b""
    i = 0
    lineprog = re.compile(r"^\s*([A-Z]+)\s*=\s*(([0-9a-f])*)\s*$", re.IGNORECASE)
    for line in test_file.readlines():
        line = line.strip()
        res = lineprog.match(line)
        if line == "":
            print()
            print("Count = %d" % i)
            print("   m = %s" % m.hex())
            print("  ad = %s" % ad.hex())
            print("npub = %s" % npub.hex())
            print("   k = %s" % k.hex())
            print("   c = %s" % c.hex())

            submit('c', b"\0" * (len(m) + 32))
            submit('s', b"")

            submit('m', m)
            submit('a', ad)
            submit('k', k)
            submit('p', npub)
            write(b'e')
            write(b'C')
            output = obtain()
            print("   c = %s" % output.hex())
            if c != output:
                raise Exception("output of encryption is different from expected ciphertext")

            submit('m', b"\0" * len(c))
            submit('s', b"")

            submit('c', c)
            submit('a', ad)
            submit('k', k)
            submit('p', npub)
            write(b'd')
            write(b'M')
            output = obtain()
            print("   m = %s" % output.hex())
            if m != output:
                raise Exception("output of encryption is different from expected ciphertext")

        elif res is not None:
            if res[1].lower() == 'count':
                i = int(res[2], 10)
            elif res[1].lower() == 'key':
                k = bytes.fromhex(res[2])
            elif res[1].lower() == 'nonce':
                npub = bytes.fromhex(res[2])
            elif res[1].lower() == 'pt':
                m = bytes.fromhex(res[2])
            elif res[1].lower() == 'ad':
                ad = bytes.fromhex(res[2])
            elif res[1].lower() == 'ct':
                c = bytes.fromhex(res[2])
            else:
Enrico Pozzobon committed
132
                raise Exception("ERROR: unparsed line in test vectors file: '%s'" % res)
Enrico Pozzobon committed
133
        else:
Enrico Pozzobon committed
134
            raise Exception("ERROR: unparsed line in test vectors file: '%s'" % line)
Enrico Pozzobon committed
135

136 137 138 139
def begin_measurement():
    import saleae
    import time
    sal = saleae.Saleae()
140
    sal.set_active_channels([0, 1, 2, 3], [])
141
    sal.set_sample_rate(sal.get_all_sample_rates()[0])
142
    sal.set_capture_seconds(1000)
143 144
    sal.capture_start()
    time.sleep(1)
145 146
    if sal.is_processing_complete():
        raise Exception("Capture didn't start successfully")
147 148 149 150 151 152
    return sal

def end_measurement(sal):
    import time
    time.sleep(1)
    sal.capture_stop();
153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169
    for attempt in range(3):
        if not sal.is_processing_complete():
            print("waiting for capture to complete")
            time.sleep(1)
            continue
        outfile = "measurement_%s.csv" % time.strftime("%Y%m%d-%H%M%S")
        outfile = os.path.join("measurements", outfile)
        if os.path.isfile(outfile):
            os.unlink(outfile)
        sal.export_data2(os.path.abspath(outfile))
        print("Measurements written to '%s'" % outfile)
        mdbfile = os.path.join("measurements", "measurements.txt")
        mdbfile = open(mdbfile, "a")
        mdbfile.write("%s > %s\n" % (' '.join(sys.argv), outfile))
        mdbfile.close()
        return 0
    raise Exception("Capture didn't complete successfully")
Enrico Pozzobon committed
170 171 172

if __name__ == "__main__":
    sys.exit(main(sys.argv))