| 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) |
