WhoDat/file_parser.py at master · applesauce777/WhoDat · GitHub

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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
import pandas as pd
import re
from pathlib import Path

def detect_ip_column(df):
    """Automatically detect the column containing IP addresses"""
    ip_pattern = re.compile(r'^(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)$')

    # Score each column based on how many valid IPs it contains
    column_scores = {}

    for column in df.columns:
        score = 0
        total_non_empty = 0

        for value in df[column].dropna():
            if pd.notna(value):
                total_non_empty += 1
                # Convert to string and check if it's an IP
                str_value = str(value).strip()
                if ip_pattern.match(str_value):
                    score += 1

        # Calculate percentage of valid IPs
        if total_non_empty > 0:
            percentage = (score / total_non_empty) * 100
            column_scores[column] = percentage

    # Return the column with the highest percentage of valid IPs
    if column_scores:
        best_column = max(column_scores, key=column_scores.get)
        if column_scores[best_column] > 50:  # At least 50% should be valid IPs
            return best_column, column_scores[best_column]

    return None, 0

def extract_ips_from_file(file_path):
    """Extract IP addresses from various file formats"""
    file_path = Path(file_path)
    file_extension = file_path.suffix.lower()

    try:
        if file_extension == '.txt':
            # Handle plain text file (one IP per line)
            with open(file_path, 'r') as f:
                ips = [line.strip() for line in f if line.strip()]
            return ips, 'text', None

        elif file_extension in ['.csv', '.tsv']:
            # Handle CSV/TSV files
            delimiter = '\t' if file_extension == '.tsv' else ','
            df = pd.read_csv(file_path, delimiter=delimiter)

            ip_column, confidence = detect_ip_column(df)
            if ip_column:
                ips = [str(ip).strip() for ip in df[ip_column].dropna() if str(ip).strip()]
                return ips, f'csv/tsv (column: {ip_column})', confidence
            else:
                return [], 'csv/tsv', 0

        elif file_extension in ['.xlsx', '.xls']:
            # Handle Excel files
            df = pd.read_excel(file_path)

            ip_column, confidence = detect_ip_column(df)
            if ip_column:
                ips = [str(ip).strip() for ip in df[ip_column].dropna() if str(ip).strip()]
                return ips, f'excel (column: {ip_column})', confidence
            else:
                return [], 'excel', 0

        else:
            return [], f'unsupported format: {file_extension}', 0

    except Exception as e:
        return [], f'error: {str(e)}', 0

def preview_file(file_path, max_rows=10):
    """Preview file contents to help users understand the structure"""
    file_path = Path(file_path)
    file_extension = file_path.suffix.lower()

    try:
        if file_extension == '.txt':
            with open(file_path, 'r') as f:
                lines = f.readlines()[:max_rows]
            return {
                'type': 'text',
                'preview': lines,
                'total_lines': len(open(file_path, 'r').readlines())
            }

        elif file_extension in ['.csv', '.tsv']:
            delimiter = '\t' if file_extension == '.tsv' else ','
            df = pd.read_csv(file_path, delimiter=delimiter)

            ip_column, confidence = detect_ip_column(df)

            return {
                'type': 'csv/tsv',
                'preview': df.head(max_rows).to_dict('records'),
                'columns': list(df.columns),
                'total_rows': len(df),
                'ip_column': ip_column,
                'confidence': confidence
            }

        elif file_extension in ['.xlsx', '.xls']:
            df = pd.read_excel(file_path)

            ip_column, confidence = detect_ip_column(df)

            return {
                'type': 'excel',
                'preview': df.head(max_rows).to_dict('records'),
                'columns': list(df.columns),
                'total_rows': len(df),
                'ip_column': ip_column,
                'confidence': confidence
            }

        else:
            return {'type': 'unsupported', 'error': f'Unsupported file format: {file_extension}'}

    except Exception as e:
        return {'type': 'error', 'error': str(e)}