use base64::prelude::*;
use serde::{Deserialize, Serialize};
use std::{fs, io::Write};
use toml::{self, Table, Value};
#[derive(Debug)]
pub struct AuthData {
username: String,
api_key: String,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ConfigFile {
auth: AuthSection,
jira: JiraSection,
#[serde(default)]
yara: YaraSection,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct AuthSection {
auth_token: String,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct JiraSection {
jira_url: String,
standard_resolution: String,
standard_resolution_comment: String,
transitions_names: Table,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct YaraSection {
rules_source: String,
rules_directory: String,
cache_file: String,
cache_version_file: String,
}
impl AuthData {
pub fn new(username: String, api_key: String) -> AuthData {
AuthData { username, api_key }
}
pub fn set_username(&mut self, username: String) {
self.username = username.replace("\n", "");
}
pub fn set_api_key(&mut self, api_key: String) {
self.api_key = api_key.replace("\n", "");
}
pub fn to_base64(&self) -> String {
BASE64_STANDARD.encode(format!("{}:{}", self.username, self.api_key).replace("\n", ""))
}
pub fn from_base64(encoded: &str) -> (String, String) {
let decoded = BASE64_STANDARD
.decode(encoded)
.expect("Failed to decode base64");
let decoded_str = String::from_utf8(decoded).expect("Failed to convert to string");
let parts: Vec<&str> = decoded_str.split(':').collect();
(parts[0].to_string(), parts[1].to_string())
}
}
impl ConfigFile {
pub fn new(
auth_token: String,
jira_url: String,
standard_resolution: String,
standard_resolution_comment: String,
transitions_names: Table,
yara: YaraSection,
) -> ConfigFile {
ConfigFile {
auth: AuthSection { auth_token },
jira: JiraSection {
jira_url,
standard_resolution,
standard_resolution_comment,
transitions_names,
},
yara,
}
}
pub fn set_auth_key(&mut self, auth_token: String) {
self.auth.auth_token = auth_token.replace("\n", "");
}
pub fn get_auth_key(&self) -> &str {
&self.auth.auth_token
}
pub fn set_jira_url(&mut self, jira_url: String) {
self.jira.jira_url = jira_url.replace("\n", "");
}
pub fn get_jira_url(&self) -> &str {
&self.jira.jira_url
}
pub fn set_standard_resolution(&mut self, standard_resolution: String) {
self.jira.standard_resolution = standard_resolution;
}
pub fn get_standard_resolution(&self) -> &String {
&self.jira.standard_resolution
}
pub fn set_standard_resolution_comment(&mut self, standard_resolution_comment: String) {
self.jira.standard_resolution_comment = standard_resolution_comment;
}
pub fn get_standard_resolution_comment(&self) -> &String {
&self.jira.standard_resolution_comment
}
pub fn add_transition_name(&mut self, key: String, value: String) {
let mut existing_value: Vec<Value> = self
.jira
.transitions_names
.get(&key)
.and_then(|v| v.as_array())
.unwrap_or(&vec![])
.iter()
.map(|v| Value::String(v.as_str().unwrap().to_string()))
.collect();
existing_value.push(Value::String(value));
self.jira
.transitions_names
.insert(key, Value::Array(existing_value));
}
pub fn get_transition_name(&self, key: &str) -> Option<Vec<String>> {
let tranisitons_names = self
.jira
.transitions_names
.get(key)
.and_then(|v| v.as_array());
Some(
tranisitons_names
.unwrap_or(&vec![])
.iter()
.map(|v| v.as_str().unwrap().to_string())
.collect(),
)
}
pub fn get_yara_section(&self) -> &YaraSection {
&self.yara
}
pub fn get_yara_rules_source(&self) -> &str {
&self.yara.rules_source
}
pub fn get_yara_rules_directory(&self) -> &str {
&self.yara.rules_directory
}
pub fn get_yara_cache_file(&self) -> &str {
&self.yara.cache_file
}
pub fn get_yara_cache_version_file(&self) -> &str {
&self.yara.cache_version_file
}
pub fn set_yara_rules_source(&mut self, source: String) {
self.yara.rules_source = source;
}
pub fn set_yara_rules_directory(&mut self, directory: String) {
self.yara.rules_directory = directory;
}
pub fn set_yara_cache_file(&mut self, file: String) {
self.yara.cache_file = file;
}
pub fn set_yara_cache_version_file(&mut self, file: String) {
self.yara.cache_version_file = file;
}
pub fn write_to_file(&self, file_path: &str) -> Result<(), std::io::Error> {
let mut file = std::fs::OpenOptions::new()
.write(true)
.create(true)
.truncate(true)
.open(file_path)
.expect("Failed to open file");
let toml_str = toml::to_string(self).expect("Failed to serialize toml");
file.write_all(toml_str.as_bytes())
}
pub fn read_from_file(file_path: &str) -> Result<ConfigFile, toml::de::Error> {
let config_file_str = fs::read_to_string(file_path)
.unwrap_or(toml::to_string(&ConfigFile::default()).unwrap_or("".to_string()));
toml::from_str(&config_file_str)
}
}
impl Default for YaraSection {
fn default() -> YaraSection {
YaraSection {
rules_source: String::from(
"https://github.com/YARAHQ/yara-forge/releases/latest/download/yara-forge-rules-core.zip",
),
rules_directory: String::from("yara-rules"),
cache_file: String::from("yara_rules.cache"),
cache_version_file: String::from("yara_rules.cache.version"),
}
}
}
impl Default for ConfigFile {
fn default() -> ConfigFile {
ConfigFile {
auth: AuthSection {
auth_token: String::from(""),
},
jira: JiraSection {
jira_url: String::from(""),
standard_resolution: String::from(""),
standard_resolution_comment: String::from(""),
transitions_names: Table::new(),
},
yara: YaraSection::default(),
}
}
}
impl YaraSection {
pub fn new(
rules_source: String,
rules_directory: String,
cache_file: String,
cache_version_file: String,
) -> YaraSection {
YaraSection {
rules_source,
rules_directory,
cache_file,
cache_version_file,
}
}
pub fn get_rules_source(&self) -> &str {
&self.rules_source
}
pub fn get_rules_directory(&self) -> &str {
&self.rules_directory
}
pub fn get_cache_file(&self) -> &str {
&self.cache_file
}
pub fn get_cache_version_file(&self) -> &str {
&self.cache_version_file
}
}