PyScan/functionScanner.py at main · S4Plus/PyScan · 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
113
114
115
116
117
118
119
120
121
122
123
124
import ast
import getopt
import sys
import csv
from bs4 import BeautifulSoup
import re

class scanner(ast.NodeVisitor):
    def __init__(self, html, parametic_poly):
        self.funcs = {}
        self.funcs["funcs"] = []
        self.funcs["classes"] = {}
        self.classnames = []
        self.parametic_poly = parametic_poly
        if html != None:
            self.soup = BeautifulSoup(html, features="html.parser")

    def visit_FunctionDef(self, node):
        if len(self.classnames) == 0 and not node.name.startswith("_") and node.name not in self.funcs["funcs"] and self.parametic_poly == True and self.check_polymorphism(node.name, node.lineno):
            self.funcs["funcs"].append(node.name)
        elif len(self.classnames) == 1 and not node.name.startswith("_") and node.name not in self.funcs["classes"][self.classnames[len(self.classnames)  - 1]] and self.parametic_poly == True and self.check_polymorphism(node.name, node.lineno):
            self.funcs["classes"][self.classnames[len(self.classnames)  - 1]].append(node.name)
        elif len(self.classnames) == 0 and not node.name.startswith("_") and node.name not in self.funcs["funcs"] and self.parametic_poly == False:
            self.funcs["funcs"].append(node.name)
        elif len(self.classnames) == 1 and not node.name.startswith("_") and node.name not in self.funcs["classes"][self.classnames[len(self.classnames)  - 1]] and self.parametic_poly == False:
            self.funcs["classes"][self.classnames[len(self.classnames)  - 1]].append(node.name)

    def visit_ClassDef(self, node):
        self.classnames.append(node.name)

        if self.classnames[len(self.classnames) - 1] not in self.funcs["classes"]:
            self.funcs["classes"][self.classnames[len(self.classnames) - 1]] = []

        self.generic_visit(node)

        self.classnames.pop(len(self.classnames) - 1)

    def finalize(self):
        deletekeys = []
        for c in self.funcs["classes"]:
            if len(self.funcs["classes"][c]) == 0:
                deletekeys.append(c)
        for c in deletekeys:
            del self.funcs["classes"][c]

    def print_to_csv(self, csvfile):
        keys = ["class", "func"]
        values = []
        for func in self.funcs["funcs"]:
            values.append(["None", func])
        for c in self.funcs["classes"]:
            for func in self.funcs["classes"][c]:
                values.append([c, func])

        if len(values) > 0:
            with open(csvfile, "w") as csvcontent:
                writer = csv.writer(csvcontent)
                writer.writerow(keys)
                writer.writerows(values)

    def check_polymorphism(self, name, lineno):
        spans = self.soup.find_all("span", text = re.compile(" *" + str(lineno)))
        if len(spans) > 1 or len(spans) == 0:
            return False
        else:
            p = spans[0].next_sibling
            while(p != None and p.name != "span"):
                if hasattr(p, "xid") and hasattr(p, "title") and p.string == name:
                    types = p["title"].split(" / ")
                    count = 0
                    returnvalues = {}
                    for t in types:
                        if "?" not in t and " -> " in t:
                            sig = t.split(" -> ")
                            if "|" in sig[0] and "{" in sig[0] and "}" in sig[0]:
                                return True
                            if sig[1] not in returnvalues:
                                returnvalues[sig[1]] = sig[0]
                            elif returnvalues[sig[1]] != sig[0]:
                                count += 1
                    if count > 0:
                        return True
                    else:
                        return False
                p = p.next_sibling
        return False


if __name__ == "__main__":
    sourcefile = None
    csvfile = None
    parametic_poly = False
    htmlfile = None
    try:
        opts, args = getopt.getopt(sys.argv[1:],"-h-s:-c:-p-t:",["source=", "csvfile=", "parametic-poly", "typeres="])
    except getopt.GetoptError:
        print("Unsupportable arguments, please see functionScanner.py -h")
        sys.exit(-1)
    for opt, arg in opts:
        if opt == '-h':
            print("Usage:\n-s/--source <Python Source File> : Indicate the path of Python source file")
            print("-c/--csvfile <CSV File> : Write the result into the csv file")
            sys.exit()
        elif opt in ("-s", "--source"):
            sourcefile = arg
        elif opt in ("-c", "--csvfile"):
            csvfile = arg
        elif opt in ("-p", "--parametic-poly"):
            parametic_poly = True
        elif opt in ("-t", "--typeres"):
            htmlfile = arg

    if sourcefile != None:
        source = open(sourcefile, "r").read()
        root = ast.parse(source)
        html = open(htmlfile, "r")
        visitor = scanner(html, parametic_poly)
        visitor.visit(root)
        visitor.finalize()
        print(visitor.funcs)
        if csvfile != None:
            visitor.print_to_csv(csvfile)