June1作品--赛题1 #9
|
|
@ -0,0 +1,50 @@
|
|||
FROM continuumio/miniconda3:4.9.2
|
||||
|
||||
ENV PYTHONUNBUFFERED=1 \
|
||||
PYTHONDONTWRITEBYTECODE=1 \
|
||||
PIP_NO_CACHE_DIR=off \
|
||||
PIP_DISABLE_PIP_VERSION_CHECK=on \
|
||||
PIP_DEFAULT_TIMEOUT=100 \
|
||||
POETRY_VERSION=1.1.11 \
|
||||
PYTHONPATH="/app:$PYTHONPATH"
|
||||
|
||||
RUN apt-get update && apt-get install -y \
|
||||
build-essential \
|
||||
cmake \
|
||||
git \
|
||||
curl \
|
||||
vim \
|
||||
wget \
|
||||
libssl-dev \
|
||||
zlib1g-dev \
|
||||
libbz2-dev \
|
||||
libreadline-dev \
|
||||
libsqlite3-dev \
|
||||
libncursesw5-dev \
|
||||
xz-utils \
|
||||
tk-dev \
|
||||
libxml2-dev \
|
||||
libxmlsec1-dev \
|
||||
libffi-dev \
|
||||
liblzma-dev \
|
||||
mecab-ipadic-utf8
|
||||
|
||||
RUN conda create -n cvss_env python=3.8 -y
|
||||
|
||||
SHELL ["conda", "run", "-n", "cvss_env", "/bin/bash", "-c"]
|
||||
|
||||
RUN pip install "poetry==$POETRY_VERSION"
|
||||
|
||||
WORKDIR /app
|
||||
COPY pyproject.toml poetry.lock ./
|
||||
|
||||
RUN poetry config virtualenvs.create false \
|
||||
&& poetry install --no-dev --no-interaction --no-ansi
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN pip install -r requirements.txt
|
||||
|
||||
RUN python setup.py build_ext --inplace
|
||||
|
||||
ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "cvss_env", "python", "main.py"]
|
||||
14
README.md
14
README.md
|
|
@ -1,14 +0,0 @@
|
|||
## :rocket: 背景
|
||||
|
||||
近年来,开源软件供应链遭受持续的软件投毒和恶意代码攻击,造成了无法估计的损失。例如,Apache Log4j2远程代码执行漏洞被认为是近10年最严重的漏洞之一,攻击者可以在目标服务器上执行任意代码和嗅探系统信息。网络安全专家认为Log4j 中的远程代码执行漏洞可能需要数月甚至数年时间才能得到妥善解决。受Log4J漏洞影响组件包括Apache的Struts2、Solr、Druid、Flink等,Github上60,644个开源项目发布321,094软件存在风险。因此,当前急需智能化技术辅助降低漏洞风险,提高漏洞工程能力,减少漏洞损失。
|
||||
|
||||
|
||||
## :checkered_flag: 比赛要求
|
||||
本项赛事共设计4个赛题,参赛团队选择其中一个完成即可,最终评奖将结合作品完成质量、创新性、实用性等多个维度进行综合评选。
|
||||
|
||||
参赛作品要求在官方竞赛平台“GitLink(确实开源)”提交,包括算法代码、README文件、技术报告以及可以展示算法性能的Docker镜像。具体要求参见赛事网站的“参赛指南”。
|
||||
|
||||
注:推荐使用开源大模型。
|
||||
|
||||
|
||||
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
# CVSS 评分系统
|
||||
|
||||
## 简介
|
||||
|
||||
第七届开源大赛智能漏洞检测项目作品
|
||||
|
||||
## 主要功能
|
||||
|
||||
- 自动化数据收集和预处理
|
||||
- 多模型集成学习
|
||||
- 高级特征工程
|
||||
- 自动化超参数调优
|
||||
- 可解释性分析
|
||||
- 交互式可视化报告生成
|
||||
|
||||
## 系统要求
|
||||
|
||||
- Python 3.8+
|
||||
- CUDA 兼容 GPU(推荐用于深度学习模型)
|
||||
- 最少 16GB RAM
|
||||
- 50GB 可用磁盘空间
|
||||
|
||||
## 环境
|
||||
|
||||
1.安装Miniconda
|
||||
2.新建conda环境
|
||||
3.pip install -r requirements.txt
|
||||
|
||||
## 使用方法
|
||||
|
||||
python main.py
|
||||
Binary file not shown.
|
|
@ -0,0 +1,4 @@
|
|||
from .data_ingestion import DataIngester
|
||||
from .feature_craft import FeatureCrafter
|
||||
|
||||
__all__ = ['DataIngester', 'FeatureCrafter']
|
||||
|
|
@ -0,0 +1,34 @@
|
|||
import json
|
||||
import pandas as pd
|
||||
from typing import List, Dict
|
||||
|
||||
class DataIngester:
|
||||
def __init__(self, file_paths: List[str]):
|
||||
self.file_paths = file_paths
|
||||
self.data = []
|
||||
|
||||
def load_data(self) -> pd.DataFrame:
|
||||
for file_path in self.file_paths:
|
||||
with open(file_path, 'r') as f:
|
||||
self.data.extend(json.load(f))
|
||||
return pd.DataFrame(self.data)
|
||||
|
||||
def preprocess_data(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
df['Issue_Created_At'] = pd.to_datetime(df['Issue_Created_At'])
|
||||
df['description_length'] = df['description'].str.len()
|
||||
df['has_url'] = df['description'].str.contains('http|https').astype(int)
|
||||
df['vectorString_length'] = df['vectorString'].str.len()
|
||||
return df
|
||||
|
||||
def split_data(self, df: pd.DataFrame, test_size: float = 0.2, val_size: float = 0.1) -> Dict[str, pd.DataFrame]:
|
||||
train_size = 1 - test_size - val_size
|
||||
train = df.sample(frac=train_size, random_state=42)
|
||||
remaining = df.drop(train.index)
|
||||
val = remaining.sample(frac=val_size/(test_size+val_size), random_state=42)
|
||||
test = remaining.drop(val.index)
|
||||
return {'train': train, 'val': val, 'test': test}
|
||||
|
||||
def process(self) -> Dict[str, pd.DataFrame]:
|
||||
df = self.load_data()
|
||||
df = self.preprocess_data(df)
|
||||
return self.split_data(df)
|
||||
|
|
@ -0,0 +1,30 @@
|
|||
import pandas as pd
|
||||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||||
from sklearn.preprocessing import StandardScaler
|
||||
|
||||
class FeatureCrafter:
|
||||
def __init__(self, max_features: int = 1000):
|
||||
self.tfidf = TfidfVectorizer(max_features=max_features)
|
||||
self.scaler = StandardScaler()
|
||||
|
||||
def fit_transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
text_features = self.tfidf.fit_transform(df['description'])
|
||||
text_feature_names = self.tfidf.get_feature_names_out()
|
||||
text_df = pd.DataFrame(text_features.toarray(), columns=text_feature_names)
|
||||
|
||||
numeric_features = ['description_length', 'has_url', 'vectorString_length']
|
||||
numeric_df = df[numeric_features]
|
||||
numeric_df = pd.DataFrame(self.scaler.fit_transform(numeric_df), columns=numeric_features)
|
||||
|
||||
return pd.concat([numeric_df, text_df], axis=1)
|
||||
|
||||
def transform(self, df: pd.DataFrame) -> pd.DataFrame:
|
||||
text_features = self.tfidf.transform(df['description'])
|
||||
text_feature_names = self.tfidf.get_feature_names_out()
|
||||
text_df = pd.DataFrame(text_features.toarray(), columns=text_feature_names)
|
||||
|
||||
numeric_features = ['description_length', 'has_url', 'vectorString_length']
|
||||
numeric_df = df[numeric_features]
|
||||
numeric_df = pd.DataFrame(self.scaler.transform(numeric_df), columns=numeric_features)
|
||||
|
||||
return pd.concat([numeric_df, text_df], axis=1)
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
from .model_arena import ModelArena
|
||||
from .hyper_tuner import HyperTuner
|
||||
|
||||
__all__ = ['ModelArena', 'HyperTuner']
|
||||
|
|
@ -0,0 +1,80 @@
|
|||
from sklearn.model_selection import RandomizedSearchCV
|
||||
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from typing import Dict, Any
|
||||
|
||||
class HyperTuner:
|
||||
def __init__(self, X: pd.DataFrame, y: pd.Series, n_iter: int = 100, cv: int = 5):
|
||||
self.X = X
|
||||
self.y = y
|
||||
self.n_iter = n_iter
|
||||
self.cv = cv
|
||||
|
||||
def tune_random_forest(self) -> Dict[str, Any]:
|
||||
param_dist = {
|
||||
'n_estimators': np.arange(100, 1000, 100),
|
||||
'max_depth': [None] + list(np.arange(10, 110, 10)),
|
||||
'min_samples_split': np.arange(2, 12, 2),
|
||||
'min_samples_leaf': np.arange(1, 11, 2),
|
||||
'max_features': ['auto', 'sqrt', 'log2']
|
||||
}
|
||||
rf = RandomForestRegressor(random_state=42)
|
||||
random_search = RandomizedSearchCV(rf, param_distributions=param_dist, n_iter=self.n_iter, cv=self.cv, random_state=42, n_jobs=-1)
|
||||
random_search.fit(self.X, self.y)
|
||||
return {
|
||||
'best_params': random_search.best_params_,
|
||||
'best_score': random_search.best_score_,
|
||||
'best_estimator': random_search.best_estimator_
|
||||
}
|
||||
|
||||
def tune_gradient_boosting(self) -> Dict[str, Any]:
|
||||
param_dist = {
|
||||
'n_estimators': np.arange(100, 1000, 100),
|
||||
'learning_rate': [0.01, 0.05, 0.1, 0.2],
|
||||
'max_depth': np.arange(3, 10),
|
||||
'min_samples_split': np.arange(2, 12, 2),
|
||||
'min_samples_leaf': np.arange(1, 11, 2),
|
||||
'subsample': [0.6, 0.7, 0.8, 0.9, 1.0],
|
||||
'max_features': ['auto', 'sqrt', 'log2']
|
||||
}
|
||||
gb = GradientBoostingRegressor(random_state=42)
|
||||
random_search = RandomizedSearchCV(gb, param_distributions=param_dist, n_iter=self.n_iter, cv=self.cv, random_state=42, n_jobs=-1)
|
||||
random_search.fit(self.X, self.y)
|
||||
return {
|
||||
'best_params': random_search.best_params_,
|
||||
'best_score': random_search.best_score_,
|
||||
'best_estimator': random_search.best_estimator_
|
||||
}
|
||||
|
||||
def tune_experimental_net(self) -> Dict[str, Any]:
|
||||
from model_zoo.experimental_net import ExperimentalNet
|
||||
param_dist = {
|
||||
'hidden_layers': [[64, 32], [128, 64], [256, 128, 64]],
|
||||
'dropout_rate': [0.1, 0.2, 0.3, 0.4, 0.5],
|
||||
'learning_rate': [0.001, 0.01, 0.1],
|
||||
'batch_size': [32, 64, 128],
|
||||
'epochs': [50, 100, 200]
|
||||
}
|
||||
def create_model(hidden_layers, dropout_rate, learning_rate):
|
||||
return ExperimentalNet(input_dim=self.X.shape[1], hidden_layers=hidden_layers, dropout_rate=dropout_rate, learning_rate=learning_rate)
|
||||
|
||||
best_score = float('inf')
|
||||
best_params = None
|
||||
best_model = None
|
||||
|
||||
for _ in range(self.n_iter):
|
||||
params = {k: np.random.choice(v) for k, v in param_dist.items()}
|
||||
model = create_model(params['hidden_layers'], params['dropout_rate'], params['learning_rate'])
|
||||
history = model.fit(self.X, self.y, epochs=params['epochs'], batch_size=params['batch_size'], validation_split=0.2, verbose=0)
|
||||
score = min(history.history['val_loss'])
|
||||
if score < best_score:
|
||||
best_score = score
|
||||
best_params = params
|
||||
best_model = model
|
||||
|
||||
return {
|
||||
'best_params': best_params,
|
||||
'best_score': best_score,
|
||||
'best_estimator': best_model
|
||||
}
|
||||
|
|
@ -0,0 +1,47 @@
|
|||
import pandas as pd
|
||||
import numpy as np
|
||||
from sklearn.metrics import mean_squared_error
|
||||
from typing import Dict, List
|
||||
from model_zoo.deep_learner import DeepLearner
|
||||
from model_zoo.ensemble_factory import EnsembleFactory
|
||||
from model_zoo.experimental_net import ExperimentalNet
|
||||
|
||||
class ModelArena:
|
||||
def __init__(self, X_train: pd.DataFrame, y_train: pd.Series, X_val: pd.DataFrame, y_val: pd.Series):
|
||||
self.X_train = X_train
|
||||
self.y_train = y_train
|
||||
self.X_val = X_val
|
||||
self.y_val = y_val
|
||||
self.models = self._initialize_models()
|
||||
self.results = {}
|
||||
|
||||
def _initialize_models(self) -> Dict:
|
||||
return {
|
||||
'DeepLearner': DeepLearner(),
|
||||
'RandomForest': EnsembleFactory().create_random_forest(),
|
||||
'GradientBoosting': EnsembleFactory().create_gradient_boosting(),
|
||||
'ExperimentalNet': ExperimentalNet()
|
||||
}
|
||||
|
||||
def train_and_evaluate(self):
|
||||
for name, model in self.models.items():
|
||||
model.fit(self.X_train, self.y_train)
|
||||
train_pred = model.predict(self.X_train)
|
||||
val_pred = model.predict(self.X_val)
|
||||
self.results[name] = {
|
||||
'train_mse': mean_squared_error(self.y_train, train_pred),
|
||||
'val_mse': mean_squared_error(self.y_val, val_pred),
|
||||
'train_rmse': np.sqrt(mean_squared_error(self.y_train, train_pred)),
|
||||
'val_rmse': np.sqrt(mean_squared_error(self.y_val, val_pred))
|
||||
}
|
||||
|
||||
def get_best_model(self, metric: str = 'val_rmse') -> str:
|
||||
return min(self.results, key=lambda x: self.results[x][metric])
|
||||
|
||||
def get_results_dataframe(self) -> pd.DataFrame:
|
||||
return pd.DataFrame(self.results).T
|
||||
|
||||
def ensemble_top_models(self, top_n: int = 3, metric: str = 'val_rmse') -> np.ndarray:
|
||||
top_models = sorted(self.results, key=lambda x: self.results[x][metric])[:top_n]
|
||||
predictions = np.array([self.models[model].predict(self.X_val) for model in top_models])
|
||||
return np.mean(predictions, axis=0)
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
from .plot_master import PlotMaster
|
||||
from .report_builder import ReportBuilder
|
||||
|
||||
__all__ = ['PlotMaster', 'ReportBuilder']
|
||||
|
|
@ -0,0 +1,75 @@
|
|||
import matplotlib.pyplot as plt
|
||||
import seaborn as sns
|
||||
import pandas as pd
|
||||
import numpy as np
|
||||
from typing import Dict, List
|
||||
|
||||
class PlotMaster:
|
||||
def __init__(self):
|
||||
plt.style.use('seaborn')
|
||||
sns.set_palette("deep")
|
||||
|
||||
def plot_feature_importance(self, feature_importance: Dict[str, float], top_n: int = 20):
|
||||
sorted_features = sorted(feature_importance.items(), key=lambda x: x[1], reverse=True)[:top_n]
|
||||
features, importances = zip(*sorted_features)
|
||||
|
||||
plt.figure(figsize=(12, 8))
|
||||
sns.barplot(x=list(importances), y=list(features))
|
||||
plt.title(f"Top {top_n} Feature Importances")
|
||||
plt.xlabel("Importance")
|
||||
plt.ylabel("Features")
|
||||
plt.tight_layout()
|
||||
return plt
|
||||
|
||||
def plot_prediction_vs_actual(self, y_true: np.ndarray, y_pred: np.ndarray):
|
||||
plt.figure(figsize=(10, 10))
|
||||
plt.scatter(y_true, y_pred, alpha=0.5)
|
||||
plt.plot([y_true.min(), y_true.max()], [y_true.min(), y_true.max()], 'r--', lw=2)
|
||||
plt.xlabel("Actual CVSS Score")
|
||||
plt.ylabel("Predicted CVSS Score")
|
||||
plt.title("Predicted vs Actual CVSS Scores")
|
||||
plt.tight_layout()
|
||||
return plt
|
||||
|
||||
def plot_error_distribution(self, y_true: np.ndarray, y_pred: np.ndarray):
|
||||
errors = y_pred - y_true
|
||||
plt.figure(figsize=(10, 6))
|
||||
sns.histplot(errors, kde=True)
|
||||
plt.xlabel("Prediction Error")
|
||||
plt.ylabel("Frequency")
|
||||
plt.title("Distribution of Prediction Errors")
|
||||
plt.tight_layout()
|
||||
return plt
|
||||
|
||||
def plot_model_comparison(self, results: Dict[str, Dict[str, float]]):
|
||||
df = pd.DataFrame(results).T
|
||||
metrics = ['train_rmse', 'val_rmse']
|
||||
|
||||
plt.figure(figsize=(12, 6))
|
||||
df[metrics].plot(kind='bar', width=0.8)
|
||||
plt.title("Model Performance Comparison")
|
||||
plt.xlabel("Models")
|
||||
plt.ylabel("RMSE")
|
||||
plt.legend(title="Metrics")
|
||||
plt.xticks(rotation=45)
|
||||
plt.tight_layout()
|
||||
return plt
|
||||
|
||||
def plot_learning_curve(self, train_sizes: np.ndarray, train_scores: List[float], val_scores: List[float]):
|
||||
plt.figure(figsize=(10, 6))
|
||||
plt.plot(train_sizes, np.mean(train_scores, axis=1), label='Training score')
|
||||
plt.plot(train_sizes, np.mean(val_scores, axis=1), label='Validation score')
|
||||
plt.xlabel("Training Set Size")
|
||||
plt.ylabel("Score")
|
||||
plt.title("Learning Curve")
|
||||
plt.legend(loc="best")
|
||||
plt.grid(True)
|
||||
plt.tight_layout()
|
||||
return plt
|
||||
|
||||
def plot_correlation_heatmap(self, df: pd.DataFrame):
|
||||
plt.figure(figsize=(12, 10))
|
||||
sns.heatmap(df.corr(), annot=True, cmap='coolwarm', linewidths=0.5)
|
||||
plt.title("Feature Correlation Heatmap")
|
||||
plt.tight_layout()
|
||||
return plt
|
||||
|
|
@ -0,0 +1,82 @@
|
|||
import pandas as pd
|
||||
import numpy as np
|
||||
from typing import Dict, List
|
||||
from jinja2 import Template
|
||||
import matplotlib.pyplot as plt
|
||||
import io
|
||||
import base64
|
||||
|
||||
class ReportBuilder:
|
||||
def __init__(self):
|
||||
self.report_data = {}
|
||||
|
||||
def add_model_performance(self, model_results: Dict[str, Dict[str, float]]):
|
||||
self.report_data['model_performance'] = pd.DataFrame(model_results).T
|
||||
|
||||
def add_feature_importance(self, feature_importance: Dict[str, float], top_n: int = 20):
|
||||
sorted_features = sorted(feature_importance.items(), key=lambda x: x[1], reverse=True)[:top_n]
|
||||
self.report_data['feature_importance'] = pd.DataFrame(sorted_features, columns=['Feature', 'Importance'])
|
||||
|
||||
def add_prediction_stats(self, y_true: np.ndarray, y_pred: np.ndarray):
|
||||
self.report_data['prediction_stats'] = {
|
||||
'mse': np.mean((y_true - y_pred) ** 2),
|
||||
'rmse': np.sqrt(np.mean((y_true - y_pred) ** 2)),
|
||||
'mae': np.mean(np.abs(y_true - y_pred)),
|
||||
'r2': 1 - (np.sum((y_true - y_pred) ** 2) / np.sum((y_true - np.mean(y_true)) ** 2))
|
||||
}
|
||||
|
||||
def add_plot(self, plot_func, *args, **kwargs):
|
||||
plt_obj = plot_func(*args, **kwargs)
|
||||
img = io.BytesIO()
|
||||
plt_obj.savefig(img, format='png')
|
||||
img.seek(0)
|
||||
plot_url = base64.b64encode(img.getvalue()).decode()
|
||||
if 'plots' not in self.report_data:
|
||||
self.report_data['plots'] = []
|
||||
self.report_data['plots'].append(plot_url)
|
||||
plt.close()
|
||||
|
||||
def generate_html_report(self) -> str:
|
||||
template = Template("""
|
||||
<html>
|
||||
<head>
|
||||
<title>CVSS Evaluation Report</title>
|
||||
<style>
|
||||
body { font-family: Arial, sans-serif; }
|
||||
table { border-collapse: collapse; width: 100%; }
|
||||
th, td { border: 1px solid #ddd; padding: 8px; }
|
||||
th { background-color: #f2f2f2; }
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<h1>CVSS Evaluation Report</h1>
|
||||
|
||||
<h2>Model Performance</h2>
|
||||
{{ model_performance.to_html() | safe }}
|
||||
|
||||
<h2>Feature Importance</h2>
|
||||
{{ feature_importance.to_html() | safe }}
|
||||
|
||||
<h2>Prediction Statistics</h2>
|
||||
<table>
|
||||
{% for key, value in prediction_stats.items() %}
|
||||
<tr>
|
||||
<th>{{ key }}</th>
|
||||
<td>{{ value }}</td>
|
||||
</tr>
|
||||
{% endfor %}
|
||||
</table>
|
||||
|
||||
<h2>Visualizations</h2>
|
||||
{% for plot in plots %}
|
||||
<img src="data:image/png;base64,{{ plot }}" alt="Plot">
|
||||
{% endfor %}
|
||||
</body>
|
||||
</html>
|
||||
""")
|
||||
|
||||
return template.render(**self.report_data)
|
||||
|
||||
def save_html_report(self, filename: str):
|
||||
with open(filename, 'w') as f:
|
||||
f.write(self.generate_html_report())
|
||||
|
|
@ -0,0 +1,128 @@
|
|||
import pandas as pd
|
||||
import numpy as np
|
||||
from sklearn.model_selection import train_test_split
|
||||
from typing import Dict, Any
|
||||
|
||||
from data_wizardry.data_ingestion import DataIngester
|
||||
from data_wizardry.feature_craft import FeatureCrafter
|
||||
from model_zoo.deep_learner import DeepLearner
|
||||
from model_zoo.ensemble_factory import EnsembleFactory
|
||||
from model_zoo.experimental_net import ExperimentalNet
|
||||
from scoring_magic.cvss_engine import CVSSEngine
|
||||
from scoring_magic.risk_oracle import RiskOracle
|
||||
from experiment_lab.model_arena import ModelArena
|
||||
from experiment_lab.hyper_tuner import HyperTuner
|
||||
from insight_generator.plot_master import PlotMaster
|
||||
from insight_generator.report_builder import ReportBuilder
|
||||
from utils.config_loader import ConfigLoader
|
||||
from utils.performance_tracker import PerformanceTracker
|
||||
|
||||
class CVSSEvaluator:
|
||||
def __init__(self, config_path: str = 'config.yaml'):
|
||||
self.config = ConfigLoader(config_path)
|
||||
self.performance_tracker = PerformanceTracker()
|
||||
self.plot_master = PlotMaster()
|
||||
self.report_builder = ReportBuilder()
|
||||
self.cvss_engine = CVSSEngine()
|
||||
self.risk_oracle = RiskOracle()
|
||||
|
||||
@performance_tracker.track_execution_time
|
||||
def load_and_preprocess_data(self) -> Dict[str, pd.DataFrame]:
|
||||
data_ingester = DataIngester(self.config.get('data_paths'))
|
||||
data = data_ingester.process()
|
||||
|
||||
feature_crafter = FeatureCrafter()
|
||||
processed_data = {}
|
||||
for split, df in data.items():
|
||||
X = feature_crafter.fit_transform(df) if split == 'train' else feature_crafter.transform(df)
|
||||
y = df['baseScore']
|
||||
processed_data[split] = {'X': X, 'y': y}
|
||||
|
||||
return processed_data
|
||||
|
||||
@performance_tracker.track_execution_time
|
||||
def train_models(self, data: Dict[str, Dict[str, pd.DataFrame]]) -> Dict[str, Any]:
|
||||
X_train, y_train = data['train']['X'], data['train']['y']
|
||||
X_val, y_val = data['val']['X'], data['val']['y']
|
||||
|
||||
model_arena = ModelArena(X_train, y_train, X_val, y_val)
|
||||
model_arena.train_and_evaluate()
|
||||
|
||||
best_model_name = model_arena.get_best_model()
|
||||
best_model = model_arena.models[best_model_name]
|
||||
|
||||
return {
|
||||
'best_model': best_model,
|
||||
'best_model_name': best_model_name,
|
||||
'all_results': model_arena.get_results_dataframe()
|
||||
}
|
||||
|
||||
@performance_tracker.track_execution_time
|
||||
def hyperparameter_tuning(self, data: Dict[str, Dict[str, pd.DataFrame]]) -> Dict[str, Any]:
|
||||
X_train, y_train = data['train']['X'], data['train']['y']
|
||||
hyper_tuner = HyperTuner(X_train, y_train)
|
||||
|
||||
rf_results = hyper_tuner.tune_random_forest()
|
||||
gb_results = hyper_tuner.tune_gradient_boosting()
|
||||
exp_results = hyper_tuner.tune_experimental_net()
|
||||
|
||||
return {
|
||||
'random_forest': rf_results,
|
||||
'gradient_boosting': gb_results,
|
||||
'experimental_net': exp_results
|
||||
}
|
||||
|
||||
@performance_tracker.track_execution_time
|
||||
def evaluate_model(self, model: Any, data: Dict[str, Dict[str, pd.DataFrame]]) -> Dict[str, float]:
|
||||
X_test, y_test = data['test']['X'], data['test']['y']
|
||||
y_pred = model.predict(X_test)
|
||||
|
||||
self.performance_tracker.calculate_metrics(y_test, y_pred)
|
||||
return self.performance_tracker.get_metrics()
|
||||
|
||||
def generate_visualizations(self, data: Dict[str, Dict[str, pd.DataFrame]], model: Any, results: pd.DataFrame):
|
||||
X_test, y_test = data['test']['X'], data['test']['y']
|
||||
y_pred = model.predict(X_test)
|
||||
|
||||
self.plot_master.plot_feature_importance(model.get_feature_importance(X_test))
|
||||
self.plot_master.plot_prediction_vs_actual(y_test, y_pred)
|
||||
self.plot_master.plot_error_distribution(y_test, y_pred)
|
||||
self.plot_master.plot_model_comparison(results)
|
||||
|
||||
def generate_report(self, model: Any, data: Dict[str, Dict[str, pd.DataFrame]], results: pd.DataFrame, tuning_results: Dict[str, Any]):
|
||||
self.report_builder.add_model_performance(results.to_dict())
|
||||
self.report_builder.add_feature_importance(model.get_feature_importance(data['test']['X']))
|
||||
self.report_builder.add_prediction_stats(data['test']['y'], model.predict(data['test']['X']))
|
||||
|
||||
for plot_func in [self.plot_master.plot_feature_importance,
|
||||
self.plot_master.plot_prediction_vs_actual,
|
||||
self.plot_master.plot_error_distribution,
|
||||
self.plot_master.plot_model_comparison]:
|
||||
self.report_builder.add_plot(plot_func, *plot_func.__code__.co_varnames[1:])
|
||||
|
||||
return self.report_builder.generate_html_report()
|
||||
|
||||
def run(self):
|
||||
data = self.load_and_preprocess_data()
|
||||
model_results = self.train_models(data)
|
||||
best_model = model_results['best_model']
|
||||
all_results = model_results['all_results']
|
||||
|
||||
tuning_results = self.hyperparameter_tuning(data)
|
||||
|
||||
evaluation_metrics = self.evaluate_model(best_model, data)
|
||||
self.generate_visualizations(data, best_model, all_results)
|
||||
|
||||
report = self.generate_report(best_model, data, all_results, tuning_results)
|
||||
|
||||
print(f"Best Model: {model_results['best_model_name']}")
|
||||
print(f"Evaluation Metrics: {evaluation_metrics}")
|
||||
print(f"Performance Report:\n{self.performance_tracker.generate_performance_report()}")
|
||||
print(f"Full Report saved as 'cvss_evaluation_report.html'")
|
||||
|
||||
with open('cvss_evaluation_report.html', 'w') as f:
|
||||
f.write(report)
|
||||
|
||||
if __name__ == "__main__":
|
||||
evaluator = CVSSEvaluator()
|
||||
evaluator.run()
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
from .deep_learner import DeepLearner
|
||||
from .ensemble_factory import EnsembleFactory
|
||||
from .experimental_net import ExperimentalNet
|
||||
|
||||
__all__ = ['DeepLearner', 'EnsembleFactory', 'ExperimentalNet']
|
||||
|
|
@ -0,0 +1,41 @@
|
|||
import tensorflow as tf
|
||||
from tensorflow.keras.models import Sequential
|
||||
from tensorflow.keras.layers import Dense, Dropout
|
||||
from tensorflow.keras.optimizers import Adam
|
||||
from tensorflow.keras.callbacks import EarlyStopping
|
||||
import numpy as np
|
||||
|
||||
class DeepLearner:
|
||||
def __init__(self, input_dim, hidden_layers=[64, 32], dropout_rate=0.2, learning_rate=0.001):
|
||||
self.model = self._build_model(input_dim, hidden_layers, dropout_rate, learning_rate)
|
||||
|
||||
def _build_model(self, input_dim, hidden_layers, dropout_rate, learning_rate):
|
||||
model = Sequential()
|
||||
model.add(Dense(hidden_layers[0], activation='relu', input_dim=input_dim))
|
||||
model.add(Dropout(dropout_rate))
|
||||
|
||||
for units in hidden_layers[1:]:
|
||||
model.add(Dense(units, activation='relu'))
|
||||
model.add(Dropout(dropout_rate))
|
||||
|
||||
model.add(Dense(1))
|
||||
model.compile(optimizer=Adam(learning_rate=learning_rate), loss='mse')
|
||||
return model
|
||||
|
||||
def fit(self, X, y, epochs=100, batch_size=32, validation_split=0.2, verbose=0):
|
||||
early_stopping = EarlyStopping(patience=10, restore_best_weights=True)
|
||||
return self.model.fit(X, y, epochs=epochs, batch_size=batch_size,
|
||||
validation_split=validation_split,
|
||||
callbacks=[early_stopping], verbose=verbose)
|
||||
|
||||
def predict(self, X):
|
||||
return self.model.predict(X).flatten()
|
||||
|
||||
def get_feature_importance(self, X):
|
||||
input_tensor = tf.convert_to_tensor(X, dtype=tf.float32)
|
||||
with tf.GradientTape() as tape:
|
||||
tape.watch(input_tensor)
|
||||
predictions = self.model(input_tensor)
|
||||
gradients = tape.gradient(predictions, input_tensor)
|
||||
feature_importance = np.mean(np.abs(gradients.numpy()), axis=0)
|
||||
return feature_importance / np.sum(feature_importance)
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
|
||||
import numpy as np
|
||||
|
||||
class EnsembleFactory:
|
||||
@staticmethod
|
||||
def create_random_forest(n_estimators=100, max_depth=None, min_samples_split=2, min_samples_leaf=1):
|
||||
return RandomForestRegressor(
|
||||
n_estimators=n_estimators,
|
||||
max_depth=max_depth,
|
||||
min_samples_split=min_samples_split,
|
||||
min_samples_leaf=min_samples_leaf,
|
||||
random_state=42
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def create_gradient_boosting(n_estimators=100, learning_rate=0.1, max_depth=3, min_samples_split=2, min_samples_leaf=1):
|
||||
return GradientBoostingRegressor(
|
||||
n_estimators=n_estimators,
|
||||
learning_rate=learning_rate,
|
||||
max_depth=max_depth,
|
||||
min_samples_split=min_samples_split,
|
||||
min_samples_leaf=min_samples_leaf,
|
||||
random_state=42
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def create_stacking_ensemble(base_models, meta_model):
|
||||
from sklearn.ensemble import StackingRegressor
|
||||
return StackingRegressor(
|
||||
estimators=base_models,
|
||||
final_estimator=meta_model,
|
||||
cv=5
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def create_voting_ensemble(models, weights=None):
|
||||
from sklearn.ensemble import VotingRegressor
|
||||
return VotingRegressor(
|
||||
estimators=[(f"model_{i}", model) for i, model in enumerate(models)],
|
||||
weights=weights
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def create_bagging_ensemble(base_estimator, n_estimators=10):
|
||||
from sklearn.ensemble import BaggingRegressor
|
||||
return BaggingRegressor(
|
||||
base_estimator=base_estimator,
|
||||
n_estimators=n_estimators,
|
||||
random_state=42
|
||||
)
|
||||
|
|
@ -0,0 +1,55 @@
|
|||
import tensorflow as tf
|
||||
from tensorflow.keras.models import Model
|
||||
from tensorflow.keras.layers import Input, Dense, Dropout, BatchNormalization, Concatenate
|
||||
from tensorflow.keras.optimizers import Adam
|
||||
from tensorflow.keras.callbacks import EarlyStopping
|
||||
import numpy as np
|
||||
|
||||
class ExperimentalNet:
|
||||
def __init__(self, input_dim, hidden_layers=[64, 32], dropout_rate=0.2, learning_rate=0.001):
|
||||
self.model = self._build_model(input_dim, hidden_layers, dropout_rate, learning_rate)
|
||||
|
||||
def _build_model(self, input_dim, hidden_layers, dropout_rate, learning_rate):
|
||||
inputs = Input(shape=(input_dim,))
|
||||
|
||||
# Main branch
|
||||
x = Dense(hidden_layers[0], activation='relu')(inputs)
|
||||
x = BatchNormalization()(x)
|
||||
x = Dropout(dropout_rate)(x)
|
||||
|
||||
for units in hidden_layers[1:]:
|
||||
x = Dense(units, activation='relu')(x)
|
||||
x = BatchNormalization()(x)
|
||||
x = Dropout(dropout_rate)(x)
|
||||
|
||||
# Residual branch
|
||||
residual = Dense(hidden_layers[-1], activation='relu')(inputs)
|
||||
residual = BatchNormalization()(residual)
|
||||
|
||||
# Combine main and residual branches
|
||||
combined = Concatenate()([x, residual])
|
||||
|
||||
# Output layer
|
||||
outputs = Dense(1)(combined)
|
||||
|
||||
model = Model(inputs=inputs, outputs=outputs)
|
||||
model.compile(optimizer=Adam(learning_rate=learning_rate), loss='mse')
|
||||
return model
|
||||
|
||||
def fit(self, X, y, epochs=100, batch_size=32, validation_split=0.2, verbose=0):
|
||||
early_stopping = EarlyStopping(patience=10, restore_best_weights=True)
|
||||
return self.model.fit(X, y, epochs=epochs, batch_size=batch_size,
|
||||
validation_split=validation_split,
|
||||
callbacks=[early_stopping], verbose=verbose)
|
||||
|
||||
def predict(self, X):
|
||||
return self.model.predict(X).flatten()
|
||||
|
||||
def get_feature_importance(self, X):
|
||||
input_tensor = tf.convert_to_tensor(X, dtype=tf.float32)
|
||||
with tf.GradientTape() as tape:
|
||||
tape.watch(input_tensor)
|
||||
predictions = self.model(input_tensor)
|
||||
gradients = tape.gradient(predictions, input_tensor)
|
||||
feature_importance = np.mean(np.abs(gradients.numpy()), axis=0)
|
||||
return feature_importance / np.sum(feature_importance)
|
||||
|
|
@ -0,0 +1,41 @@
|
|||
numpy==1.21.5
|
||||
pandas==1.3.5
|
||||
scikit-learn==0.24.2
|
||||
tensorflow==2.6.0
|
||||
torch==1.10.0
|
||||
xgboost==1.5.0
|
||||
lightgbm==3.3.2
|
||||
catboost==1.0.3
|
||||
dask==2021.12.0
|
||||
vaex==4.5.0
|
||||
pyarrow==6.0.0
|
||||
matplotlib==3.5.1
|
||||
seaborn==0.11.2
|
||||
plotly==5.4.0
|
||||
bokeh==2.4.2
|
||||
statsmodels==0.13.1
|
||||
scipy==1.7.3
|
||||
sympy==1.9
|
||||
keras==2.6.0
|
||||
pytorch-lightning==1.5.0
|
||||
transformers==4.12.5
|
||||
nltk==3.6.5
|
||||
spacy==3.2.0
|
||||
gensim==4.1.2
|
||||
optuna==2.10.0
|
||||
hyperopt==0.2.5
|
||||
shap==0.40.0
|
||||
lime==0.2.0.1
|
||||
tqdm==4.62.3
|
||||
joblib==1.1.0
|
||||
dill==0.3.4
|
||||
flask==2.0.2
|
||||
fastapi==0.70.0
|
||||
pytest==6.2.5
|
||||
hypothesis==6.24.0
|
||||
sphinx==4.3.2
|
||||
mkdocs==1.2.3
|
||||
networkx==2.6.3
|
||||
numba==0.54.1
|
||||
cython==0.29.24
|
||||
pyyaml==6.0
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
from .cvss_engine import CVSSEngine
|
||||
from .risk_oracle import RiskOracle
|
||||
|
||||
__all__ = ['CVSSEngine', 'RiskOracle']
|
||||
|
|
@ -0,0 +1,121 @@
|
|||
import numpy as np
|
||||
from typing import Dict, List
|
||||
|
||||
class CVSSEngine:
|
||||
def __init__(self):
|
||||
self.vector_weights = {
|
||||
'AV': {'N': 0.85, 'A': 0.62, 'L': 0.55, 'P': 0.2},
|
||||
'AC': {'L': 0.77, 'H': 0.44},
|
||||
'PR': {'N': 0.85, 'L': 0.62, 'H': 0.27},
|
||||
'UI': {'N': 0.85, 'R': 0.62},
|
||||
'S': {'U': 6.42, 'C': 7.52},
|
||||
'C': {'N': 0, 'L': 0.22, 'H': 0.56},
|
||||
'I': {'N': 0, 'L': 0.22, 'H': 0.56},
|
||||
'A': {'N': 0, 'L': 0.22, 'H': 0.56}
|
||||
}
|
||||
|
||||
def parse_vector(self, vector_string: str) -> Dict[str, str]:
|
||||
vector_parts = vector_string.split('/')
|
||||
vector_dict = {}
|
||||
for part in vector_parts:
|
||||
if ':' in part:
|
||||
key, value = part.split(':')
|
||||
vector_dict[key] = value
|
||||
return vector_dict
|
||||
|
||||
def calculate_base_score(self, vector_dict: Dict[str, str]) -> float:
|
||||
impact_sub_score = self._calculate_impact_sub_score(vector_dict)
|
||||
exploitability_sub_score = self._calculate_exploitability_sub_score(vector_dict)
|
||||
|
||||
if vector_dict['S'] == 'U':
|
||||
base_score = min((impact_sub_score + exploitability_sub_score), 10)
|
||||
else:
|
||||
base_score = min(1.08 * (impact_sub_score + exploitability_sub_score), 10)
|
||||
|
||||
return round(base_score, 1)
|
||||
|
||||
def _calculate_impact_sub_score(self, vector_dict: Dict[str, str]) -> float:
|
||||
c = self.vector_weights['C'][vector_dict['C']]
|
||||
i = self.vector_weights['I'][vector_dict['I']]
|
||||
a = self.vector_weights['A'][vector_dict['A']]
|
||||
|
||||
if vector_dict['S'] == 'U':
|
||||
return 6.42 * (1 - (1 - c) * (1 - i) * (1 - a))
|
||||
else:
|
||||
return 7.52 * (1 - (1 - c) * (1 - i) * (1 - a)) - 3.25 * (1 - (1 - c) * (1 - i) * (1 - a)) ** 15
|
||||
|
||||
def _calculate_exploitability_sub_score(self, vector_dict: Dict[str, str]) -> float:
|
||||
av = self.vector_weights['AV'][vector_dict['AV']]
|
||||
ac = self.vector_weights['AC'][vector_dict['AC']]
|
||||
pr = self.vector_weights['PR'][vector_dict['PR']]
|
||||
ui = self.vector_weights['UI'][vector_dict['UI']]
|
||||
|
||||
return 8.22 * av * ac * pr * ui
|
||||
|
||||
def calculate_temporal_score(self, base_score: float, vector_dict: Dict[str, str]) -> float:
|
||||
e = self._get_temporal_weight(vector_dict, 'E')
|
||||
rl = self._get_temporal_weight(vector_dict, 'RL')
|
||||
rc = self._get_temporal_weight(vector_dict, 'RC')
|
||||
|
||||
return round(base_score * e * rl * rc, 1)
|
||||
|
||||
def _get_temporal_weight(self, vector_dict: Dict[str, str], metric: str) -> float:
|
||||
temporal_weights = {
|
||||
'E': {'X': 1, 'H': 1, 'F': 0.97, 'P': 0.94, 'U': 0.91},
|
||||
'RL': {'X': 1, 'U': 1, 'W': 0.97, 'T': 0.96, 'O': 0.95},
|
||||
'RC': {'X': 1, 'C': 1, 'R': 0.96, 'U': 0.92}
|
||||
}
|
||||
return temporal_weights[metric].get(vector_dict.get(metric, 'X'), 1)
|
||||
|
||||
def calculate_environmental_score(self, vector_dict: Dict[str, str]) -> float:
|
||||
modified_vector = self._apply_environmental_modifiers(vector_dict)
|
||||
modified_impact_sub_score = self._calculate_modified_impact_sub_score(modified_vector)
|
||||
modified_exploitability_sub_score = self._calculate_exploitability_sub_score(modified_vector)
|
||||
|
||||
if modified_vector['MS'] == 'U':
|
||||
environmental_score = min((modified_impact_sub_score + modified_exploitability_sub_score), 10)
|
||||
else:
|
||||
environmental_score = min(1.08 * (modified_impact_sub_score + modified_exploitability_sub_score), 10)
|
||||
|
||||
return round(environmental_score, 1)
|
||||
|
||||
def _apply_environmental_modifiers(self, vector_dict: Dict[str, str]) -> Dict[str, str]:
|
||||
modified_vector = vector_dict.copy()
|
||||
for metric in ['C', 'I', 'A']:
|
||||
if f'M{metric}' in vector_dict:
|
||||
modified_vector[metric] = vector_dict[f'M{metric}']
|
||||
for metric in ['AV', 'AC', 'PR', 'UI']:
|
||||
if f'M{metric}' in vector_dict:
|
||||
modified_vector[metric] = vector_dict[f'M{metric}']
|
||||
modified_vector['MS'] = vector_dict.get('MS', vector_dict['S'])
|
||||
return modified_vector
|
||||
|
||||
def _calculate_modified_impact_sub_score(self, vector_dict: Dict[str, str]) -> float:
|
||||
mc = self.vector_weights['C'][vector_dict['C']] * self._get_environmental_weight(vector_dict, 'CR')
|
||||
mi = self.vector_weights['I'][vector_dict['I']] * self._get_environmental_weight(vector_dict, 'IR')
|
||||
ma = self.vector_weights['A'][vector_dict['A']] * self._get_environmental_weight(vector_dict, 'AR')
|
||||
|
||||
if vector_dict['MS'] == 'U':
|
||||
return 6.42 * (1 - (1 - mc) * (1 - mi) * (1 - ma))
|
||||
else:
|
||||
return 7.52 * (1 - (1 - mc) * (1 - mi) * (1 - ma)) - 3.25 * (1 - (1 - mc) * (1 - mi) * (1 - ma)) ** 15
|
||||
|
||||
def _get_environmental_weight(self, vector_dict: Dict[str, str], metric: str) -> float:
|
||||
environmental_weights = {
|
||||
'CR': {'X': 1, 'H': 1.5, 'M': 1, 'L': 0.5},
|
||||
'IR': {'X': 1, 'H': 1.5, 'M': 1, 'L': 0.5},
|
||||
'AR': {'X': 1, 'H': 1.5, 'M': 1, 'L': 0.5}
|
||||
}
|
||||
return environmental_weights[metric].get(vector_dict.get(metric, 'X'), 1)
|
||||
|
||||
def calculate_overall_score(self, vector_string: str) -> Dict[str, float]:
|
||||
vector_dict = self.parse_vector(vector_string)
|
||||
base_score = self.calculate_base_score(vector_dict)
|
||||
temporal_score = self.calculate_temporal_score(base_score, vector_dict)
|
||||
environmental_score = self.calculate_environmental_score(vector_dict)
|
||||
|
||||
return {
|
||||
'base_score': base_score,
|
||||
'temporal_score': temporal_score,
|
||||
'environmental_score': environmental_score
|
||||
}
|
||||
|
|
@ -0,0 +1,80 @@
|
|||
import numpy as np
|
||||
from typing import Dict, List
|
||||
from .cvss_engine import CVSSEngine
|
||||
|
||||
class RiskOracle:
|
||||
def __init__(self):
|
||||
self.cvss_engine = CVSSEngine()
|
||||
self.risk_levels = {
|
||||
'Critical': (9.0, 10.0),
|
||||
'High': (7.0, 8.9),
|
||||
'Medium': (4.0, 6.9),
|
||||
'Low': (0.1, 3.9),
|
||||
'None': (0.0, 0.0)
|
||||
}
|
||||
|
||||
def assess_risk(self, vector_string: str) -> Dict[str, any]:
|
||||
scores = self.cvss_engine.calculate_overall_score(vector_string)
|
||||
base_risk = self._get_risk_level(scores['base_score'])
|
||||
temporal_risk = self._get_risk_level(scores['temporal_score'])
|
||||
environmental_risk = self._get_risk_level(scores['environmental_score'])
|
||||
|
||||
return {
|
||||
'base_risk': base_risk,
|
||||
'temporal_risk': temporal_risk,
|
||||
'environmental_risk': environmental_risk,
|
||||
'scores': scores
|
||||
}
|
||||
|
||||
def _get_risk_level(self, score: float) -> str:
|
||||
for level, (min_score, max_score) in self.risk_levels.items():
|
||||
if min_score <= score <= max_score:
|
||||
return level
|
||||
return 'Unknown'
|
||||
|
||||
def generate_risk_report(self, vector_string: str) -> str:
|
||||
assessment = self.assess_risk(vector_string)
|
||||
report = f"CVSS Vector: {vector_string}\n\n"
|
||||
report += "Risk Assessment:\n"
|
||||
report += f"Base Risk: {assessment['base_risk']} (Score: {assessment['scores']['base_score']})\n"
|
||||
report += f"Temporal Risk: {assessment['temporal_risk']} (Score: {assessment['scores']['temporal_score']})\n"
|
||||
report += f"Environmental Risk: {assessment['environmental_risk']} (Score: {assessment['scores']['environmental_score']})\n\n"
|
||||
report += self._generate_recommendations(assessment)
|
||||
return report
|
||||
|
||||
def _generate_recommendations(self, assessment: Dict[str, any]) -> str:
|
||||
recommendations = "Recommendations:\n"
|
||||
if assessment['environmental_risk'] in ['Critical', 'High']:
|
||||
recommendations += "- Immediate action required to mitigate the vulnerability\n"
|
||||
recommendations += "- Implement temporary workarounds if a patch is not immediately available\n"
|
||||
recommendations += "- Conduct a thorough impact analysis\n"
|
||||
elif assessment['environmental_risk'] == 'Medium':
|
||||
recommendations += "- Develop and implement a remediation plan within a reasonable timeframe\n"
|
||||
recommendations += "- Prioritize based on business impact and exploitation likelihood\n"
|
||||
else:
|
||||
recommendations += "- Address the vulnerability as part of regular maintenance cycles\n"
|
||||
recommendations += "- Monitor for any changes in exploit availability or impact\n"
|
||||
|
||||
recommendations += "- Regularly reassess the vulnerability as new information becomes available\n"
|
||||
return recommendations
|
||||
|
||||
def batch_assess_risk(self, vector_strings: List[str]) -> List[Dict[str, any]]:
|
||||
return [self.assess_risk(vector) for vector in vector_strings]
|
||||
|
||||
def get_risk_distribution(self, vector_strings: List[str]) -> Dict[str, int]:
|
||||
assessments = self.batch_assess_risk(vector_strings)
|
||||
distribution = {level: 0 for level in self.risk_levels.keys()}
|
||||
for assessment in assessments:
|
||||
distribution[assessment['environmental_risk']] += 1
|
||||
return distribution
|
||||
|
||||
def prioritize_vulnerabilities(self, vector_strings: List[str]) -> List[Dict[str, any]]:
|
||||
assessments = self.batch_assess_risk(vector_strings)
|
||||
prioritized = sorted(
|
||||
assessments,
|
||||
key=lambda x: (
|
||||
list(self.risk_levels.keys()).index(x['environmental_risk']),
|
||||
-x['scores']['environmental_score']
|
||||
)
|
||||
)
|
||||
return prioritized
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
from .config_loader import ConfigLoader
|
||||
from .performance_tracker import PerformanceTracker
|
||||
|
||||
__all__ = ['ConfigLoader', 'PerformanceTracker']
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
import yaml
|
||||
import os
|
||||
from typing import Dict, Any
|
||||
|
||||
class ConfigLoader:
|
||||
def __init__(self, config_path: str = 'config.yaml'):
|
||||
self.config_path = config_path
|
||||
self.config = self._load_config()
|
||||
|
||||
def _load_config(self) -> Dict[str, Any]:
|
||||
if not os.path.exists(self.config_path):
|
||||
raise FileNotFoundError(f"Config file not found: {self.config_path}")
|
||||
|
||||
with open(self.config_path, 'r') as config_file:
|
||||
try:
|
||||
return yaml.safe_load(config_file)
|
||||
except yaml.YAMLError as e:
|
||||
raise ValueError(f"Error parsing config file: {e}")
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
return self.config.get(key, default)
|
||||
|
||||
def get_nested(self, *keys: str, default: Any = None) -> Any:
|
||||
value = self.config
|
||||
for key in keys:
|
||||
if isinstance(value, dict):
|
||||
value = value.get(key, default)
|
||||
else:
|
||||
return default
|
||||
return value
|
||||
|
||||
def update_config(self, new_config: Dict[str, Any]):
|
||||
self.config.update(new_config)
|
||||
with open(self.config_path, 'w') as config_file:
|
||||
yaml.dump(self.config, config_file)
|
||||
|
||||
def get_all(self) -> Dict[str, Any]:
|
||||
return self.config.copy()
|
||||
|
||||
def validate_config(self, required_keys: List[str]) -> bool:
|
||||
for key in required_keys:
|
||||
if key not in self.config:
|
||||
raise KeyError(f"Required key '{key}' not found in config")
|
||||
return True
|
||||
|
||||
def get_model_config(self, model_name: str) -> Dict[str, Any]:
|
||||
model_config = self.get_nested('models', model_name)
|
||||
if model_config is None:
|
||||
raise KeyError(f"Configuration for model '{model_name}' not found")
|
||||
return model_config
|
||||
|
|
@ -0,0 +1,91 @@
|
|||
import time
|
||||
from typing import Dict, List, Callable
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
class PerformanceTracker:
|
||||
def __init__(self):
|
||||
self.metrics = {}
|
||||
self.execution_times = {}
|
||||
|
||||
def track_execution_time(self, func: Callable) -> Callable:
|
||||
def wrapper(*args, **kwargs):
|
||||
start_time = time.time()
|
||||
result = func(*args, **kwargs)
|
||||
end_time = time.time()
|
||||
execution_time = end_time - start_time
|
||||
self.execution_times[func.__name__] = execution_time
|
||||
return result
|
||||
return wrapper
|
||||
|
||||
def calculate_metrics(self, y_true: np.ndarray, y_pred: np.ndarray):
|
||||
self.metrics['mse'] = mean_squared_error(y_true, y_pred)
|
||||
self.metrics['rmse'] = np.sqrt(self.metrics['mse'])
|
||||
self.metrics['mae'] = mean_absolute_error(y_true, y_pred)
|
||||
self.metrics['r2'] = r2_score(y_true, y_pred)
|
||||
|
||||
def add_custom_metric(self, name: str, value: float):
|
||||
self.metrics[name] = value
|
||||
|
||||
def get_metrics(self) -> Dict[str, float]:
|
||||
return self.metrics
|
||||
|
||||
def get_execution_times(self) -> Dict[str, float]:
|
||||
return self.execution_times
|
||||
|
||||
def generate_performance_report(self) -> str:
|
||||
report = "Performance Report\n"
|
||||
report += "==================\n\n"
|
||||
|
||||
report += "Metrics:\n"
|
||||
for metric, value in self.metrics.items():
|
||||
report += f" {metric}: {value:.4f}\n"
|
||||
|
||||
report += "\nExecution Times:\n"
|
||||
for func, time in self.execution_times.items():
|
||||
report += f" {func}: {time:.4f} seconds\n"
|
||||
|
||||
return report
|
||||
|
||||
def plot_metrics_comparison(self, other_metrics: Dict[str, float], title: str = "Metrics Comparison"):
|
||||
metrics = list(self.metrics.keys())
|
||||
current_values = [self.metrics[m] for m in metrics]
|
||||
other_values = [other_metrics.get(m, 0) for m in metrics]
|
||||
|
||||
x = np.arange(len(metrics))
|
||||
width = 0.35
|
||||
|
||||
fig, ax = plt.subplots(figsize=(10, 6))
|
||||
ax.bar(x - width/2, current_values, width, label='Current')
|
||||
ax.bar(x + width/2, other_values, width, label='Other')
|
||||
|
||||
ax.set_ylabel('Values')
|
||||
ax.set_title(title)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(metrics)
|
||||
ax.legend()
|
||||
|
||||
plt.tight_layout()
|
||||
return plt
|
||||
|
||||
def export_metrics_to_csv(self, filename: str):
|
||||
df = pd.DataFrame([self.metrics])
|
||||
df.to_csv(filename, index=False)
|
||||
|
||||
def import_metrics_from_csv(self, filename: str):
|
||||
df = pd.read_csv(filename)
|
||||
self.metrics = df.to_dict('records')[0]
|
||||
|
||||
def track_memory_usage(self, func: Callable) -> Callable:
|
||||
import tracemalloc
|
||||
def wrapper(*args, **kwargs):
|
||||
tracemalloc.start()
|
||||
result = func(*args, **kwargs)
|
||||
current, peak = tracemalloc.get_traced_memory()
|
||||
tracemalloc.stop()
|
||||
self.metrics[f'{func.__name__}_memory_current'] = current / 10**6 # MB
|
||||
self.metrics[f'{func.__name__}_memory_peak'] = peak / 10**6 # MB
|
||||
return result
|
||||
return wrapper
|
||||
Loading…
Reference in New Issue