File size: 3,568 Bytes
77eacb7 cb8b12f 77eacb7 cb8b12f 7969312 cb8b12f 77eacb7 |
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 |
from transformers import PretrainedConfig
import torch
class CybersecurityKnowledgeGraphConfig(PretrainedConfig):
def __init__(
self,
event_nugget_model_path : str = "nugget_model_state_dict.pth",
event_argument_model_path : str = "argument_model_state_dict.pth",
event_realis_model_path : str = "realis_model_state_dict.pth",
**kwargs,
):
self.event_nugget_model_path = event_nugget_model_path
self.event_argument_model_path = event_argument_model_path
self.event_realis_model_path = event_realis_model_path
# self.event_nugget_list = utils.event_nugget_list
# self.event_args_list = utils.event_args_list
# self.realis_list = utils.realis_list
# self.arg_2_role = utils.arg_2_role
self.event_nugget_list = ['O',
'B-Ransom',
'I-Ransom',
'B-DiscoverVulnerability',
'I-DiscoverVulnerability',
'B-PatchVulnerability',
'I-PatchVulnerability',
'B-Databreach',
'I-Databreach',
'B-Phishing',
'I-Phishing'
]
self.event_args_list = ['O',
'B-System',
'I-System',
'B-Organization',
'B-Money',
'I-Money',
'B-Device',
'B-Person',
'I-Person',
'B-Vulnerability',
'I-Vulnerability',
'B-Capabilities',
'I-Capabilities',
'I-Organization',
'B-PaymentMethod',
'I-PaymentMethod',
'B-Data',
'I-Data',
'B-Number',
'I-Number',
'B-Malware',
'I-Malware',
'B-PII',
'I-PII',
'B-CVE',
'I-CVE',
'B-Purpose',
'I-Purpose',
'B-File',
'I-File',
'I-Device',
'B-Time',
'I-Time',
'B-Software',
'I-Software',
'B-Patch',
'I-Patch',
'B-Version',
'I-Version',
'B-Website',
'I-Website',
'B-GPE',
'I-GPE'
]
self.realis_list =["O",
"Generic",
"Other",
"Actual"
]
self.arg_2_role = {
"File" : ['Tool', 'Trusted-Entity'],
"Person" : ['Victim', 'Attacker', 'Discoverer', 'Releaser', 'Trusted-Entity', 'Vulnerable_System_Owner'],
"Capabilities" : ['Attack-Pattern', 'Capabilities', 'Issues-Addressed'],
"Purpose" : ['Purpose'],
"Time" : ['Time'],
"PII" : ['Compromised-Data', 'Trusted-Entity'],
"Data" : ['Compromised-Data', 'Trusted-Entity'],
"Organization" : ['Victim', 'Releaser', 'Discoverer', 'Attacker', 'Vulnerable_System_Owner', 'Trusted-Entity'],
"Patch" : ['Patch'],
"Software" : ['Vulnerable_System', 'Victim', 'Trusted-Entity', 'Supported_Platform'],
"Vulnerability" : ['Vulnerability'],
"Version" : ['Patch-Number', 'Vulnerable_System_Version'],
"Device" : ['Vulnerable_System', 'Victim', 'Supported_Platform'],
"CVE" : ['CVE'],
"Number" : ['Number-of-Data', 'Number-of-Victim'],
"System" : ['Victim', 'Supported_Platform', 'Vulnerable_System', 'Trusted-Entity'],
"Malware" : ['Tool'],
"Money" : ['Price', 'Damage-Amount'],
"PaymentMethod" : ['Payment-Method'],
"GPE" : ['Place'],
"Website" : ['Trusted-Entity', 'Tool', 'Vulnerable_System', 'Victim', 'Supported_Platform'],
}
super().__init__(**kwargs) |