"""Loads the dataset from a CSV file."""

try:

df = pd.read_csv(dataset_path)

logging.info(f"Dataset loaded successfully from {dataset_path} with {len(df)} samples.")

return df

except FileNotFoundError:

logging.error(f"Dataset not found at {dataset_path}")

return None

except Exception as e:

logging.error(f"Error loading dataset: {e}")

"""

if df is None or df.empty:

logging.warning("Preprocessing skipped: Input DataFrame is empty or None.")

return None, [], []

# Ensure target column exists and drop rows with missing target

if TARGET_COLUMN not in df.columns:

logging.error(f"Target column '{TARGET_COLUMN}' not found in the dataset.")

return None, [], []

df.dropna(subset=[TARGET_COLUMN], inplace=True)

if df.empty:

logging.warning("Preprocessing stopped: No data remaining after dropping NaNs in target column.")

return None, [], []

df[TARGET_COLUMN] = df[TARGET_COLUMN].astype(int)

# Identify initial feature columns (excluding target and potential file identifiers)

potential_id_cols = ['FileName', 'filename', 'ID', 'id'] # Add other potential ID columns if needed

feature_cols_initial = [col for col in df.columns if col != TARGET_COLUMN and col not in potential_id_cols]

# Separate numeric and categorical features

numeric_cols = df[feature_cols_initial].select_dtypes(include=np.number).columns.tolist()

categorical_cols = [col for col in feature_cols_initial if col not in numeric_cols]

logging.info(f"Identified {len(numeric_cols)} numeric and {len(categorical_cols)} categorical features.")

# Impute NaNs in numeric columns with the mean

if numeric_cols:

numeric_means = df[numeric_cols].mean()

df[numeric_cols] = df[numeric_cols].fillna(numeric_means)

nan_check = df[numeric_cols].isnull().sum().sum()

if nan_check == 0:

logging.info("NaNs imputed in numeric columns using mean.")

else:

logging.warning(f"NaN imputation failed for some numeric columns ({nan_check} NaNs remain).")

# One-hot encode categorical columns

if categorical_cols:

df = pd.get_dummies(df, columns=categorical_cols, drop_first=True, dummy_na=False)

logging.info(f"Categorical columns one-hot encoded. New shape: {df.shape}")

# Define final feature columns *after* encoding

feature_cols_final = [col for col in df.columns if col != TARGET_COLUMN and col not in potential_id_cols]

"""

scaler = None

if not numeric_cols:

logging.info("No numeric columns provided for scaling.")

return df, None

# Ensure numeric columns actually exist in the dataframe after potential prior processing steps

numeric_cols_present = [col for col in numeric_cols if col in df.columns]

if not numeric_cols_present:

logging.warning("None of the specified numeric columns found in the DataFrame for scaling.")

return df, None

try:

scaler = StandardScaler()

# Create a copy before scaling to avoid SettingWithCopyWarning

df_scaled = df.copy()

df_scaled[numeric_cols_present] = scaler.fit_transform(df[numeric_cols_present])

logging.info(f"Scaled {len(numeric_cols_present)} numeric columns.")

return df_scaled, scaler

except Exception as e:

logging.error(f"Error during feature scaling: {e}")

"""Trains multiple classification models."""

models = {

'Logistic Regression': LogisticRegression(solver='liblinear', random_state=42, max_iter=1000),

'Random Forest': RandomForestClassifier(n_estimators=100, random_state=42),

'Gradient Boosting': GradientBoostingClassifier(n_estimators=100, learning_rate=0.1, max_depth=3,

random_state=42),

'SVM': SVC(kernel='rbf', probability=True, random_state=42) # Note: SVM can be slow

}

trained_models = {}

for name, model in models.items():

logging.info(f"Training {name}...")

try:

model.fit(X_train, y_train)

trained_models[name] = model

logging.info(f"{name} trained successfully.")

except Exception as e:

logging.error(f"Error training {name}: {e}")

"""Evaluates trained models and logs metrics."""

logging.info(f"--- Model Evaluation on {dataset_label} Set ---")

evaluation_results = {}

for name, model in trained_models.items():

try:

y_pred = model.predict(X_data)

accuracy = accuracy_score(y_data, y_pred)

report = classification_report(y_data, y_pred, zero_division=0,

output_dict=True) # Get dict for easier parsing

confusion = confusion_matrix(y_data, y_pred).tolist() # Convert to list for easier logging/storage

logging.info(f"\n--- {name} ({dataset_label}) ---")

logging.info(f"Accuracy: {accuracy:.4f}")

logging.info(

f"Classification Report:\n{classification_report(y_data, y_pred, zero_division=0)}") # Log formatted

# string too

logging.info(f"Confusion Matrix:\n{confusion}")

evaluation_results[name] = {'accuracy': accuracy, 'report': report, 'confusion_matrix': confusion}

except Exception as e:

logging.error(f"Error evaluating {name} on {dataset_label} set: {e}")

"""Selects the best model based on accuracy on the test set."""

best_model_name = None

best_accuracy = -1.0

best_model = None

logging.info("--- Selecting Best Model based on Test Set Accuracy ---")

if not trained_models:

logging.warning("Cannot select best model: No models were trained successfully.")

return None, None

for name, model in trained_models.items():

try:

y_pred = model.predict(X_test)

accuracy = accuracy_score(y_test, y_pred)

logging.info(f"{name} Test Accuracy: {accuracy:.4f}")

if accuracy > best_accuracy:

best_accuracy = accuracy

best_model_name = name

best_model = model

except Exception as e:

logging.warning(f"Could not evaluate {name} for selection: {e}")

if best_model:

logging.info(f"Selected best model: {best_model_name} with accuracy {best_accuracy:.4f}")

else:

logging.warning("Could not select a best model.")

"""Saves the model, scaler, and column lists."""

try:

joblib.dump(model, MODEL_SAVE_PATH)

logging.info(f"Best model saved to: {MODEL_SAVE_PATH}")

if scaler:

joblib.dump(scaler, SCALER_SAVE_PATH)

logging.info(f"Scaler saved to: {SCALER_SAVE_PATH}")

else:

logging.info("No scaler object to save (likely no numeric features were scaled).")

# Optionally remove old scaler file if it exists

if os.path.exists(SCALER_SAVE_PATH):

os.remove(SCALER_SAVE_PATH)

with open(FEATURE_COLUMNS_PATH, 'w') as f:

for col in feature_columns:

f.write(f"{col}\n")

logging.info(f"Feature columns saved to: {FEATURE_COLUMNS_PATH}")

with open(NUMERIC_COLUMNS_PATH, 'w') as f:

for col in numeric_columns:

f.write(f"{col}\n")

logging.info(f"Numeric columns (scaled) saved to: {NUMERIC_COLUMNS_PATH}")

except Exception as e:

"""Loads the model, scaler, feature column list, and numeric column list."""

loaded_model, loaded_scaler, loaded_feature_columns, loaded_numeric_columns = None, None, None, None

try:

if os.path.exists(MODEL_SAVE_PATH):

loaded_model = joblib.load(MODEL_SAVE_PATH)

logging.info(f"Model loaded from: {MODEL_SAVE_PATH}")

else:

logging.error(f"Model file not found at {MODEL_SAVE_PATH}. Cannot proceed.")

return None, None, None, None # Model is essential

if os.path.exists(SCALER_SAVE_PATH):

loaded_scaler = joblib.load(SCALER_SAVE_PATH)

logging.info(f"Scaler loaded from: {SCALER_SAVE_PATH}")

else:

logging.info("Scaler file not found (likely no numeric features were scaled).")

if os.path.exists(FEATURE_COLUMNS_PATH):

with open(FEATURE_COLUMNS_PATH, 'r') as f:

loaded_feature_columns = [line.strip() for line in f if line.strip()]

logging.info(f"Feature columns loaded from: {FEATURE_COLUMNS_PATH}")

else:

logging.error(f"Feature columns file not found at {FEATURE_COLUMNS_PATH}. Cannot proceed with prediction.")

# Clean up potentially loaded model if feature columns are missing

if loaded_model: loaded_model = None

if loaded_scaler: loaded_scaler = None

return None, None, None, None # Feature columns are essential

if os.path.exists(NUMERIC_COLUMNS_PATH):

with open(NUMERIC_COLUMNS_PATH, 'r') as f:

loaded_numeric_columns = [line.strip() for line in f if line.strip()]

logging.info(f"Numeric columns loaded from: {NUMERIC_COLUMNS_PATH}")

else:

logging.info("Numeric columns file not found (likely no numeric features were scaled).")

loaded_numeric_columns = []

return loaded_model, loaded_scaler, loaded_feature_columns, loaded_numeric_columns

except Exception as e:

logging.error(f"Error loading components: {e}")

"""

features = {}

try:

if not os.path.exists(file_path):

logging.error(f'File not found: {file_path}')

return None

pdf_reader = PyPDF2.PdfReader(file_path)

pdfsize = os.path.getsize(file_path) / 1024 # size in KB

num_pages = len(pdf_reader.pages)

metadata = pdf_reader.metadata

metadata_size = len(str(metadata)) if metadata else 0

# Basic features matching typical column names from datasets

features['pdfsize'] = pdfsize

features['metadata size'] = metadata_size # Ensure key matches training column name

features['pages'] = num_pages

# Add more features based on your training data, e.g.:

# features['title characters'] = len(metadata.title) if metadata and metadata.title else 0

# Content-based features

js_found, acroform_found, embeddedfile_found = 0, 0, 0

try:

# Check trailer first

trailer_str = str(pdf_reader.trailer).lower() # Use lower for case-insensitivity

if '/acroform' in trailer_str: acroform_found = 1

if '/embeddedfile' in trailer_str: embeddedfile_found = 1

if '/js' in trailer_str or '/javascript' in trailer_str: js_found = 1

# Optionally search text content (can be slow and error-prone)

if not js_found:

text_content = ""

for page in pdf_reader.pages:

try:

page_text = page.extract_text()

if page_text: text_content += page_text

except Exception:

# logging.debug(f"Could not extract text from a page in {file_path}")

pass

if re.search(r'[/<\(\[\s](js|javascript)[/>\)\]\s]', text_content, re.IGNORECASE):

js_found = 1

except Exception as e:

logging.warning(f"Could not fully extract content features from {file_path}: {e}")

features['JS'] = js_found # Ensure key matches training column name

features['Acroform'] = acroform_found # Ensure key matches training column name

features['EmbeddedFile'] = embeddedfile_found # Ensure key matches training column name

logging.debug(f"Extracted features for {os.path.basename(file_path)}: {features}")

return features

except FileNotFoundError: # Already handled above, but keep for safety

logging.error(f"File not found during PDF processing: {file_path}")

return None

except PyPDF2.errors.PdfReadError as pdf_err:

logging.error(f"PyPDF2 error processing PDF {file_path}: {pdf_err}")

return None

except Exception as e:

logging.error(f"Unexpected error processing PDF {file_path}: {e}")

"""

if not raw_features_dict or not loaded_feature_columns:

logging.error("Cannot align features: Missing raw features or feature column list.")

return None

# Create a DataFrame from the single dictionary

df = pd.DataFrame([raw_features_dict], index=[0])

try:

# Reindex to match training columns: adds missing (with 0), drops extra, orders correctly.

df_aligned = df.reindex(columns=loaded_feature_columns, fill_value=0)

return df_aligned

except Exception as e:

logging.error(f"Error aligning features: {e}")

"""

if df_aligned is None or df_aligned.empty:

logging.error("Cannot scale: Input DataFrame is empty or None.")

return None

if not scaler or not numeric_columns_scaled:

logging.info("Skipping scaling: No scaler or numeric columns list provided.")

return df_aligned # Return unscaled but aligned data

# Identify numeric columns present in the aligned DataFrame that need scaling

numeric_cols_to_scale = [col for col in numeric_columns_scaled if col in df_aligned.columns]

if not numeric_cols_to_scale:

logging.info("Skipping scaling: No numeric columns from training found in current features.")

return df_aligned

try:

# Work on a copy to avoid modifying the original aligned df if scaling fails

df_scaled = df_aligned.copy()

numeric_data = df_scaled[numeric_cols_to_scale].fillna(0) # Ensure no NaNs before transform

scaled_values = scaler.transform(numeric_data)

df_scaled[numeric_cols_to_scale] = scaled_values

logging.debug("Numeric features scaled successfully for prediction.")

return df_scaled

except ValueError as ve:

logging.error(f"Scaler error during prediction (likely feature mismatch): {ve}. "

f"Columns expected by scaler: {getattr(scaler, 'n_features_in_', 'N/A')}, "

f"Columns provided: {len(numeric_cols_to_scale)}")

return None

except Exception as e:

logging.error(f"Error applying scaling during prediction: {e}")

"""

if not model or not feature_columns:

logging.error("Prediction failed: Model or feature columns not loaded.")

return {'error': 'Model or feature columns missing.'}

# 1. Extract raw features

raw_features = extract_features_single_pdf(file_path)

if raw_features is None:

return {'error': f'Failed to extract features from {os.path.basename(file_path)}.'}

# 2. Align features

aligned_df = align_features_for_prediction(raw_features, feature_columns)

if aligned_df is None:

return {'error': 'Feature alignment failed.'}

# 3. Scale features

scaled_df = scale_features_for_prediction(aligned_df, scaler, numeric_columns_scaled)

if scaled_df is None:

# Attempt prediction with unscaled data if scaling failed, but log a warning

logging.warning("Scaling failed, attempting prediction with unscaled aligned features.")

scaled_df = aligned_df # Use the aligned (but unscaled) data

# Note: This might lead to poor prediction accuracy if the model expects scaled data.

# Consider returning an error here if scaling is absolutely critical.

# return {'error': 'Feature scaling failed.'}

# 4. Predict

try:

# Ensure DataFrame columns exactly match feature_columns in order, crucial for some models

prediction_input_df = scaled_df[feature_columns]

# Final check for NaNs before prediction (should be handled by fill_value=0 and fillna(0))

if prediction_input_df.isnull().values.any():

logging.warning("NaNs found in final feature vector before prediction. Filling with 0.")

prediction_input_df = prediction_input_df.fillna(0)

prediction = model.predict(prediction_input_df)

probability = model.predict_proba(prediction_input_df)[0] if hasattr(model, 'predict_proba') else None

result = {'is_malicious': bool(prediction[0])}

if probability is not None:

# Class 1 is typically the 'malicious' class

result['malicious_probability'] = probability[1]

result['benign_probability'] = probability[0]

return result

except ValueError as e:

logging.error(f'Prediction ValueError: {e}. Check model compatibility with input shape/features. '

f'Input shape: {prediction_input_df.shape}, Expected features: {len(feature_columns)}')

return {'error': f'Prediction input shape/feature mismatch: {e}'}

except Exception as e:

logging.error(f'Prediction failed: {e}')

"""Handles the model training workflow."""

logging.info("--- Starting Training and Saving Mode ---")

df_raw = load_data(DATASET_PATH)

if df_raw is None: return # Stop if loading failed

df_processed, numeric_cols, feature_cols = preprocess_data(df_raw)

if df_processed is None or df_processed.empty: return # Stop if preprocessing failed

df_scaled, scaler = scale_features(df_processed, numeric_cols)

# Continue even if scaling fails, using df_processed

if feature_cols and TARGET_COLUMN in df_scaled.columns:

X = df_scaled[feature_cols]

y = df_scaled[TARGET_COLUMN]

if X.empty or y.empty:

logging.error("No data available for training after preprocessing/scaling.")

return

# Split data

try:

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

logging.info(f"Data split: Train set size {X_train.shape[0]}, Test set size {X_test.shape[0]}")

except ValueError as e:

logging.error(f"Error during train/test split (potentially too few samples for stratification): {e}")

return

# Train models

trained_models = train_models(X_train, y_train)

if not trained_models:

logging.error("No models were trained successfully.")

return

# Evaluate models

evaluate_models(trained_models, X_test, y_test, dataset_label="Test")

evaluate_models(trained_models, X_train, y_train, dataset_label="Training") # Evaluate on training data as well

# Select and save best model

best_model, best_model_name = select_best_model(trained_models, X_test, y_test)

if best_model:

# Use the numeric columns identified *before* scaling for saving

save_components(best_model, scaler, feature_cols, numeric_cols)

else:

logging.warning("Skipping saving components as no best model was selected.")

else:

"""Handles the prediction workflow for a single PDF."""

logging.info("--- Starting Prediction Mode ---")

# Check if PDF exists before loading components

if not os.path.exists(pdf_path_to_analyze):

logging.error(f"Input PDF file not found at '{pdf_path_to_analyze}'. Cannot predict.")

return

# Load necessary components

model, scaler, features, numeric_cols = load_components()

if model and features:

logging.info(f"Attempting prediction for: {pdf_path_to_analyze}")

prediction_result = predict_single_pdf(pdf_path_to_analyze, model, scaler, features, numeric_cols)

logging.info("-" * 30)

logging.info(f"Prediction result for {os.path.basename(pdf_path_to_analyze)}:")

if 'error' in prediction_result:

logging.error(f" Error: {prediction_result['error']}")

else:

status = "MALICIOUS" if prediction_result['is_malicious'] else "BENIGN"

logging.info(f" Status: {status}")

if 'malicious_probability' in prediction_result:

logging.info(f" Malicious Probability: {prediction_result['malicious_probability']:.4f}")

logging.info("-" * 30)

else:

logging.error("Could not load necessary components (model/feature list) for prediction. "