1#!/usr/bin/env python
2"""
3Test driver
4Runs off plain text files, similar to how PHP's test harness works
5"""
6import os
7import glob
8from libinjection import *
9from words import *
10
11print version()
12
13def print_token_string(tok):
14    """
15    returns the value of token, handling opening and closing quote characters
16    """
17    out = ''
18    if tok.str_open != "\0":
19        out += tok.str_open
20    out += tok.val
21    if tok.str_close != "\0":
22        out += tok.str_close
23    return out
24
25def print_token(tok):
26    """
27    prints a token for use in unit testing
28    """
29    out = ''
30    out += tok.type
31    out += ' '
32    if tok.type == 's':
33        out += print_token_string(tok)
34    elif tok.type == 'v':
35        vc = tok.count;
36        if vc == 1:
37            out += '@'
38        elif vc == 2:
39            out += '@@'
40        out += print_token_string(tok)
41    else:
42        out += tok.val
43    return out.strip()
44
45def toascii(data):
46    """
47    Converts a utf-8 string to ascii.   needed since nosetests xunit is not UTF-8 safe
48    https://github.com/nose-devs/nose/issues/649
49    https://github.com/nose-devs/nose/issues/692
50    """
51    return data
52    udata = data.decode('utf-8')
53    return udata.encode('ascii', 'xmlcharrefreplace')
54
55def readtestdata(filename):
56    """
57    Read a test file and split into components
58    """
59
60    state = None
61    info = {
62        '--TEST--': '',
63        '--INPUT--': '',
64        '--EXPECTED--': ''
65        }
66
67    for line in open(filename, 'r'):
68        line = line.rstrip()
69        if line in ('--TEST--', '--INPUT--', '--EXPECTED--'):
70            state = line
71        elif state:
72            info[state] += line + '\n'
73
74    # remove last newline from input
75    info['--INPUT--'] = info['--INPUT--'][0:-1]
76
77    return (info['--TEST--'], info['--INPUT--'].strip(), info['--EXPECTED--'].strip())
78
79def runtest(testname, flag, sqli_flags):
80    """
81    runs a test, optionally with valgrind
82    """
83    data =  readtestdata(os.path.join('../tests', testname))
84
85    sql_state = sqli_state()
86    sqli_init(sql_state, data[1], sqli_flags)
87    sqli_callback(sql_state, lookup)
88    actual = ''
89
90    if flag == 'tokens':
91        while sqli_tokenize(sql_state):
92            actual += print_token(sql_state.current) + '\n';
93        actual = actual.strip()
94    elif flag == 'folding':
95        num_tokens = sqli_fold(sql_state)
96        for i in range(num_tokens):
97            actual += print_token(sqli_get_token(sql_state, i)) + '\n';
98    elif flag == 'fingerprints':
99        ok = is_sqli(sql_state)
100        if ok:
101            actual = sql_state.fingerprint
102    else:
103        raise RuntimeException("unknown flag")
104
105    actual = actual.strip()
106
107    if actual != data[2]:
108        print "INPUT: \n" + toascii(data[1])
109        print
110        print "EXPECTED: \n" + toascii(data[2])
111        print
112        print "GOT: \n" + toascii(actual)
113        assert actual == data[2]
114
115def test_tokens():
116    for testname in sorted(glob.glob('../tests/test-tokens-*.txt')):
117        testname = os.path.basename(testname)
118        runtest(testname, 'tokens', libinjection.FLAG_QUOTE_NONE | libinjection.FLAG_SQL_ANSI)
119
120def test_tokens_mysql():
121    for testname in sorted(glob.glob('../tests/test-tokens_mysql-*.txt')):
122        testname = os.path.basename(testname)
123        runtest(testname, 'tokens', libinjection.FLAG_QUOTE_NONE | libinjection.FLAG_SQL_MYSQL)
124
125def test_folding():
126    for testname in sorted(glob.glob('../tests/test-folding-*.txt')):
127        testname = os.path.basename(testname)
128        runtest(testname, 'folding', libinjection.FLAG_QUOTE_NONE | libinjection.FLAG_SQL_ANSI)
129
130def test_fingerprints():
131    for testname in sorted(glob.glob('../tests/test-sqli-*.txt')):
132        testname = os.path.basename(testname)
133        runtest(testname, 'fingerprints', 0)
134
135
136if __name__ == '__main__':
137    import sys
138    sys.stderr.write("run using nosetests\n")
139    sys.exit(1)
140