pipeline.py

import pandas as pd
import re
from utils import *
from config import URL_FEATURES

class EmailFeatureExtractor:
    def __init__(self):
        self.required_features = URL_FEATURES

    def transform(self, subject: str, body: str) -> pd.DataFrame:
        # Create initial DataFrame from user input
        df = pd.DataFrame([{'subject': subject, 'body': body}])

        # 1. URL & Attachment Extraction
        df['URL'] = df.apply(extract_urls_from_body, axis=1)
        df['URL_COUNT'] = df['URL'].apply(count_urls)
        
        

        # 2. Combined Text for BERT
        df['text_combined'] = df.apply(create_combined_text, axis=1)

        # 3. IP Address Detection
        df['USE_OF_IP'] = df['URL'].apply(
            lambda x: having_ip_address(x) if x else 0
        )

        # 4. Basic URL Stats
        # Note: We apply result_type='expand' if utils returns a Series
        stats = df['URL'].apply(extract_basic_url_stats)
        df[['url_length_max', 'url_length_avg', 'url_subdom_max', 'url_subdom_avg']] = stats

        # 5. Shorteners & Suspicious Keywords
        df['short_url_count'] = df['URL'].apply(count_shortened_urls)
        df['sus_url_count'] = df['URL'].apply(suspicious_words_count)
        df['sus_url_flag'] = (df['sus_url_count'] > 0).astype(int)

        # 6. Dot Features
        df[['dot_count_max', 'dot_count_avg']] = df['URL'].apply(extract_dot_features)

        # 7. Generic Character Counts
        char_map = {'perc': '%', 'ques': '?', 'hyphen': '-', 'equal': '='}
        for name, char in char_map.items():
            df[[f'{name}_max', f'{name}_avg']] = df['URL'].apply(
                lambda x: extract_char_features(x, char, name)
            )

        return self._verify_and_order(df)

    def _verify_and_order(self, df: pd.DataFrame) -> pd.DataFrame:
        missing = [col for col in self.required_features if col not in df.columns]
        if missing:
            for col in missing:
                df[col] = 0
        
        cols_to_return = self.required_features + ['text_combined']
        return df[cols_to_return]