Spaces:
Sleeping
Sleeping
File size: 2,093 Bytes
e01340e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
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] |