AI_adversarial_interface/utils.py at master · MPSFuzz/AI_adversarial_interface · 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
125
import os

from flask import Flask, request, jsonify
from functools import wraps
import yaml

import docker

def return_0_1(code: int =200 , message: str = "done", data: dict = {"v":"k"}):
    #  this code for returning response of POST(GET) with json format
    response = {
        "code": code,
        "message": message,
        "data": data
    }

    return jsonify(response)


def share_weight(contain_id, shared_path):
    ##  this function for sharing weights of one model generated by different adversarial methods.
    return 0

def update_dict_1_level(original, new):
    # just update dict key
    for key, value in new.items():
        if key not in original:
            original[key] = value

def update_dict_2_level(original, new):
    # if new dict's model is same with origin, update the 2 level key 'weight_number', 'weight_name',
    # 'test_method', and 'download_addr'
    for key, value in new.items():
        if key not in original:
            original[key] = value
        else:
            if isinstance(original[key]['weight_number'], int):
                original[key]['weight_number'] = original[key]['weight_number'] + new[key]['weight_number']
            else:
                original[key]['weight_number'] = int(original[key]['weight_number']) + int(new[key]['weight_number'])

            for kkey in ['weight_name', 'test_method', 'download_addr']:
                if isinstance(original[key][kkey], str) and isinstance(new[key][kkey], str):
                    original[key][kkey] = [original[key][kkey], new[key][kkey]]
                elif isinstance(original[key][kkey], list) and isinstance(new[key][kkey], str):
                    original[key][kkey].append(new[key][kkey])
                elif isinstance(original[key][kkey], str) and isinstance(new[key][kkey], list):
                    original[key][kkey] = new[key][kkey].append(original[key][kkey])
                elif isinstance(original[key][kkey], list) and isinstance(new[key][kkey], list):
                    original[key][kkey].extend(new[key][kkey])


def init_read_yaml_for_model():
    yaml_file_path = './config/adver_white_box.yaml'

    with open(yaml_file_path, 'r') as yaml_file:
        data_dict = yaml.safe_load(yaml_file)

    # print(data_dict)

    new_data_dict = os.listdir("./config")
    for item in new_data_dict:
        if item != 'adver_white_box.yaml':
            # print(item)
            with open(os.path.join("./config", item)) as yaml_file:
                new_data = yaml.safe_load(yaml_file)
                update_dict_1_level(data_dict, new_data)

    # print(data_dict)
    return data_dict


def init_read_yaml_for_model_duplicate():
    yaml_file_path = './config/adver_white_box.yaml'

    with open(yaml_file_path, 'r') as yaml_file:
        data_dict = yaml.safe_load(yaml_file)

    # print(data_dict)

    new_data_dict = os.listdir("./config")
    for item in new_data_dict:
        if item != 'adver_white_box.yaml':
            with open(os.path.join("./config", item)) as yaml_file:
                new_data = yaml.safe_load(yaml_file)
                update_dict_2_level(data_dict, new_data)

    # print(data_dict)
    return data_dict

def update_yaml():
    ## do we need this funcion?  To be added
    return 0

def exec_docker_container_shell(shell_path:str) -> str:
    client = docker.from_env()

    parts = shell_path.split(":")
    container_id = parts[0]
    ###   你可以给docker容器里的shell脚本提供参数，约定好就行
    script_path = parts[1] + ' status'    ##  option: {start|stop|status|restart}

    # print("容器 ID:", container_id)
    # print("脚本路径:", script_path)
    # print("docker start %s" % (container_id))
    os.system("docker start %s" % (container_id))

    container = client.containers.get(container_id)
    exec_result = container.exec_run(cmd=script_path)


    if exec_result.exit_code == 0:
        # 将输出从bytes解码为字符串
        output = exec_result.output.decode('utf-8')
        print("Script output:", output)
    else:
        print("Script execution failed with exit code:", exec_result.exit_code)
        print("Error output:", exec_result.output.decode('utf-8'))


if __name__ == "__main__":
    data_dict = init_read_yaml_for_model()

    print(data_dict["Vgg16"]["docker_container"])
    exec_docker_container_shell(data_dict["Vgg16"]["docker_container"])