export-telegram-group-members/export.py

import argparse
import csv
import os
import sys
import tempfile
from datetime import datetime
from pathlib import Path

OUTPUT_DIR = Path('output')
CSV_FIELDS = ['id', 'username', 'first_name', 'last_name']
FORMULA_PREFIXES = ('=', '+', '-', '@', '\t', '\r')


def parse_args(argv=None):
    parser = argparse.ArgumentParser(
        description='Export a Redis-backed Telegram group crawl to CSV.',
    )
    parser.add_argument(
        'group_id',
        type=int,
        nargs='?',
        help='Telegram group ID (defaults to the first configured group)',
    )
    parser.add_argument(
        'timecrawl',
        nargs='?',
        type=parse_run_time,
        help='Crawl time in yyyymmddhhmmss format (defaults to latest)',
    )
    return parser.parse_args(argv)


def parse_run_time(value):
    if not is_valid_run_time(value):
        if len(value) != 14 or not value.isdigit():
            raise argparse.ArgumentTypeError(
                'timecrawl must use yyyymmddhhmmss format',
            )
        raise argparse.ArgumentTypeError(
            'timecrawl must be a valid date and time',
        )
    return value


def is_valid_run_time(value):
    if not isinstance(value, str) or len(value) != 14 or not value.isdigit():
        return False
    try:
        datetime.strptime(value, '%Y%m%d%H%M%S')
    except ValueError:
        return False
    return True


def main(argv=None):
    args = parse_args(argv)

    group_id = args.group_id
    try:
        # Lazy imports keep `python export.py --help` working without REDIS_URL.
        from common import get_group_export, latest_group_export_time

        if group_id is None:
            from config import load_app_config

            group_id = load_app_config()['group_ids'][0]

        run_time = args.timecrawl
        if run_time is None:
            run_time = latest_group_export_time(group_id)
            if run_time is None:
                print(f'no exports found for group {group_id}', file=sys.stderr)
                return 1
        record = get_group_export(group_id, run_time)
    except Exception as exc:
        print(f'could not read export configuration or history: {exc}', file=sys.stderr)
        return 1

    if record is None:
        print(
            f'export not found or invalid for group {group_id} at {run_time}',
            file=sys.stderr,
        )
        return 1

    try:
        members = validate_record(record, group_id, run_time)
        output_path = write_csv(group_id, run_time, members)
    except (OSError, ValueError, csv.Error) as exc:
        print(f'could not export group {group_id} at {run_time}: {exc}', file=sys.stderr)
        return 1

    print(f'exported {len(members)} members to {output_path}')
    return 0


def validate_record(record, expected_group_id, expected_time):
    if not is_valid_run_time(expected_time):
        raise ValueError('Redis export key has an invalid crawl time')
    if not isinstance(record, dict):
        raise ValueError('Redis export is not an object')
    if record.get('group_id') != expected_group_id:
        raise ValueError('Redis export group ID does not match its key')
    if record.get('time') != expected_time:
        raise ValueError('Redis export time does not match its key')

    members = record.get('members')
    if not isinstance(members, list):
        raise ValueError('Redis export members must be a list')

    validated = []
    for index, member in enumerate(members):
        if not isinstance(member, dict):
            raise ValueError(f'member {index} is not an object')
        member_id = member.get('id')
        if not isinstance(member_id, int) or isinstance(member_id, bool):
            raise ValueError(f'member {index} has an invalid ID')

        row = {'id': member_id}
        for field in CSV_FIELDS[1:]:
            value = member.get(field)
            if value is not None and not isinstance(value, str):
                raise ValueError(f'member {index} has an invalid {field}')
            row[field] = escape_formula(value)
        validated.append(row)
    return validated


def escape_formula(value):
    if value and (
        value[0] in ('\t', '\r', '\n')
        or value.lstrip().startswith(FORMULA_PREFIXES)
    ):
        return f"'{value}"
    return value


def write_csv(group_id, run_time, members, output_dir=None):
    output_dir = output_dir or OUTPUT_DIR
    output_path = output_dir / f'{group_id}-{run_time}.csv'
    output_dir.mkdir(mode=0o700, parents=True, exist_ok=True)
    output_dir.chmod(0o700)

    temp_path = None
    try:
        with tempfile.NamedTemporaryFile(
            mode='w',
            encoding='utf-8',
            newline='',
            dir=output_dir,
            prefix=f'.{output_path.name}.',
            delete=False,
        ) as csv_file:
            temp_path = Path(csv_file.name)
            writer = csv.DictWriter(csv_file, fieldnames=CSV_FIELDS)
            writer.writeheader()
            writer.writerows(members)
        os.chmod(temp_path, 0o600)
        os.replace(temp_path, output_path)
    except Exception:
        if temp_path is not None:
            temp_path.unlink(missing_ok=True)
        raise

    return output_path


if __name__ == '__main__':
    raise SystemExit(main())