saiten.py for exercise 8-1/8-2/8-3, 2013ipp1, Future University Hakodate

saiten.py

# coding: utf-8

'''
Created on 2013/06/05

@author: toshihiro
'''

import os
import os.path as P
import re
import shutil
import sys
import subprocess

TEMP_DIR = P.join(P.curdir, "TEMP")
DEPENDS_DIR = P.join(P.curdir, "bookmarkman", "src")
PACKAGE = "bookmarkman"

devnull = open('/dev/null', 'w')

PERMITTED_CLASSES = [
    "java/io/PrintStream",
    "java/lang/AssertionError",
    "java/lang/Integer",
    "java/lang/Object",
    "java/lang/String",
    "java/lang/StringBuilder",
    "java/lang/System",
    "java/util/ArrayList",
    "java/util/Arrays",
    "java/util/HashMap", 
    "java/util/HashSet",
    "java/util/Iterator",
    "java/util/Map",
    "java/util/Map$Entry",
    "java/util/Set",

    "java/util/Vector",
    "java/util/Hashtable", 
]

PERMITTED_PACKAGES = [
    "bookmarkman",
]

if not P.isdir(DEPENDS_DIR):
    sys.exit("dependencies directory not found")

def expand_source_files(jar_file_name):
    try:
        os.makedirs(P.join(TEMP_DIR, PACKAGE))
    except:
        pass

    for f in os.listdir(P.join(DEPENDS_DIR, PACKAGE)):
        if f.endswith(".java"):
            shutil.copy(P.join(DEPENDS_DIR, PACKAGE, f), P.join(TEMP_DIR, PACKAGE))
            
    target_file = P.join(TEMP_DIR, PACKAGE, "BookmarkUtility.java")
    if P.exists(target_file):
        os.remove(target_file)

    subprocess.check_call(" ".join(["unzip", jar_file_name, 
            PACKAGE + "/BookmarkUtility.java", 
            "-d", TEMP_DIR]), 
            shell=True, stdout=devnull, stderr=devnull)
    assert P.exists(target_file)

def read_author_info_from_target():
    author_line = None
    try:
        with open(P.join(TEMP_DIR, PACKAGE, "BookmarkUtility.java"), "rb") as f:
            lines = f.read().decode("utf-8").split("\n")
        for L in lines:
            L = L.rstrip()
            if L.find("@author") >= 0:
                author_line = L
    except:
        sys.exit("fail to read target file. (invalid text encoding?)")

    if not author_line:
        sys.exit("@author line not found")
        
    pat = re.compile(r"@author[\s|　]+([^\s　]+)[\s|　]+([^\s　]+)")
    m = pat.search(author_line)
    if not m:
        sys.exit("@author line invalid format")
    return m.group(1), m.group(2)

def get_target_disassembled_text():
    cur = P.abspath(P.curdir)
    os.chdir(P.join(TEMP_DIR))
    try:
        subprocess.check_call(" ".join(["javac", P.join(PACKAGE, "BookmarkUtility.java")]), 
                shell=True, stdout=devnull, stderr=devnull)
        disasm_text = subprocess.check_output(" ".join(["javap", "-c", "-p", PACKAGE + ".BookmarkUtility"]), 
                shell=True)
    finally:
        os.chdir(cur)
    
    return disasm_text

def compile_target():
    cur = P.abspath(P.curdir)
    os.chdir(P.join(TEMP_DIR))
    try:
        subprocess.check_call(" ".join(["javac", P.join(PACKAGE, "*.java")]), 
                shell=True, stdout=devnull, stderr=devnull)
    finally:
        os.chdir(cur)

def run_target(main):
    cur = P.abspath(P.curdir)
    os.chdir(P.join(TEMP_DIR))
    try:
        return subprocess.check_output(" ".join(["java", main]), 
                shell=True, stderr=subprocess.STDOUT)
    except:
        return None
    finally:
        os.chdir(cur)

def security_check_disassembled_text(disasm_text):
    pats = map(re.compile, [
        r"^(InterfaceMethod)\s+([^.]+)[.].+$", # invokeinterface
        r"^(Method)\s+([^.]+)[.].+$", # invokevirtual
        r"^(Method)\s+()[^:]+:.+$", # invokestatic
        r"^(class)\s+(.+)$", # new
        r"^(Field)\s+([^.]+)[.].+$", # getstatic
        r"^(Field)\s+()[^:]+:.+$", # getfield, putfield
        r"^(String)\s+(.+)$", # ldc
    ])
    
    referring_classes = set()
    
    lines = [L.rstrip() for L in disasm_text.split('\n')]
    for li, L in enumerate(lines):
        p = L.find("// ")
        if p >= 0:
            comment = L[p + len("// "):]
            for pat in pats:
                m = pat.match(comment)
                if m: break
            if not m:
                sys.exit("line %d: unexpected line: %s" % (li, L))
            typ = m.group(1)
            qualified_class_name = m.group(2)
            if typ in ("InterfaceMethod", "Method", "class", "Field") and qualified_class_name and \
                    not qualified_class_name.startswith('"'):
                is_permitted = False
                if qualified_class_name in PERMITTED_CLASSES:
                    is_permitted = True
                else:
                    for pp in PERMITTED_PACKAGES:
                        if qualified_class_name.startswith(pp + "/"):
                            is_permitted = True
                            break  # for pp
                if not is_permitted:
                    sys.exit("line %d: target program accessing UNPREMITTED class: %s" % (li, qualified_class_name))
                referring_classes.add(qualified_class_name)

    return sorted(referring_classes)

def main(argv):
    argv = argv[1:]
    if len(argv) == 0 or argv[0] in ("-h", "--help"):
        sys.stdout.write("usage: saiten jar_file [-v]\n")
        sys.exit(0)

    verbose = False
    if argv[-1] == "-v":
        verbose = True
        argv = argv[:-1]

    if len(argv) >= 2:
        sys.exit("too many command-line arguments")

    jar_file_name = argv[0]
    if not jar_file_name.endswith(".jar"):
        sys.exit("invalid jar file name")

    expand_source_files(jar_file_name)
    student_id = student_name = ""
    try:
        student_id, student_name = read_author_info_from_target()
    finally:
        pass  # @author not included or invalid

    disasm_text = get_target_disassembled_text()
    if verbose:
        sys.stdout.write("disassembled target:\n")
        sys.stdout.write("%s\n" % disasm_text)

    referring_classes = security_check_disassembled_text(disasm_text)
    if verbose:
        sys.stdout.write("target accessed following classes:\n")
        for c in referring_classes:
            sys.stdout.write("  %s\n" % c)

    compile_target()

    test1 = "OK" if run_target(main=PACKAGE + ".TestSet1") else "NG"
    test2 = "OK" if run_target(main=PACKAGE + ".TestSet2") else "NG"
    test3 = "OK" if run_target(main=PACKAGE + ".TestSet3") else "NG"

    sys.stdout.write(u"%s\n" % u"\t".join([u"file_name", u"student_id", u"student_name", 
            u"TestSet1", u"TestSet2", u"TestSet3"]))
    sys.stdout.write(u"%s\n" % u"\t".join([jar_file_name, student_id, student_name, 
            test1, test2, test3]))

if __name__ == '__main__':
    main(sys.argv)