|
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 = ['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) |
