python-solidity-parser/test.py at dev-litcc · litcc/python-solidity-parser · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
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
import os
import unittest
from multiprocessing import Pool
from typing import Tuple

from solidity_parser.parser import parse, objectify


def parse_test(param: Tuple[str, str]):
    address, text = param
    print(address + " start")
    try:
        kk = parse(text=text, loc=True)
        kk2 = objectify(kk)
        if kk2 is not None:
            print(address + " ok")
            file = r'tmp/check.txt'
            with open(file, 'a+') as f:
                f.write(address+"\n")
        else:
            print(address + " is none")
    except Exception as e:
        print(address + " e rr")
        raise e


class MyTestCase(unittest.TestCase):

    def test_something(self):
        """
        单独测试指定合约解析是否存在报错
        :return:
        """
        path = "tmp/test_contract"

        dir = "0xA8D5DFcf36bCb29adF9EBbD753c2978B9a982cC6"

        for files in os.listdir(path + "/" + dir):
            f = open(path + "/" + dir + "/" + files)
            contract_body = f.read()
            f.close()

            (file_name, ext) = os.path.splitext(files)

            kk = parse(text=contract_body, loc=True)
            kk2 = objectify(kk)
            # storage_info = get_storage_info(kk2, contract_name, insert_node)
            # kk = calculate_storage_key((contract_body, file_name, False))
            print(kk2)

    def test_from_file_pool(self):
        """
        批量测试文件语法解析 多进程
        :return:
        """
        path = "tmp/test_contract"

        pool = Pool(processes=16)  # 默认线程数为16

        set_vul = list()

        for dir in os.listdir(path):  # 不仅仅是文件，当前目录下的文件夹也会被认为遍历到
            # print("files", dir)

            for files in os.listdir(path + "/" + dir):
                f = open(path + "/" + dir + "/" + files)
                contract_body = f.read()
                f.close()
                (file_name, ext) = os.path.splitext(files)
                set_vul.append((dir, contract_body))

        # 创建线程池，其返回值vul_lists就是扫描结果，vul_lists为一个双层嵌套的list，每一个文件扫描到的漏洞组成一个list，每个文件文件的list组成vul_list
        vul_lists = pool.map(
            parse_test,
            set_vul
        )
        # 等待所有线程完成
        pool.close()
        pool.join()
        # self.assertEqual(True, False)  # add assertion here

    def test_from_file(self):
        """
        批量测试文件语法解析
        :return:
        """
        path = "tmp/test_contract"


        for dir in os.listdir(path):  # 不仅仅是文件，当前目录下的文件夹也会被认为遍历到
            # print("files", dir)

            for files in os.listdir(path + "/" + dir):
                f = open(path + "/" + dir + "/" + files)
                contract_body = f.read()
                f.close()
                (file_name, ext) = os.path.splitext(files)

                if os.path.exists('tmp/check.txt'):
                    f_2 = open('tmp/check.txt')
                    check_list = f_2.read()
                    f_2.close()
                    if dir not in check_list:
                        parse_test((dir, contract_body))
                else:
                    parse_test((dir, contract_body))
        # self.assertEqual(True, False)  # add assertion here


if __name__ == '__main__':
    unittest.main()