#!/usr/bin/env python3 """ S3 log file download script Used to download statement log files within a specified time range """ import boto3 import os import sys from datetime import datetime import argparse import gzip import shutil from pathlib import Path class S3LogDownloader: def __init__(self, bucket_name, prefix, output_dir="./downloads"): """ Initialize S3 downloader Args: bucket_name: S3 bucket name prefix: S3 path prefix (excluding filename) output_dir: Local download directory """ self.bucket_name = bucket_name self.prefix = prefix.rstrip('/') + '/' self.output_dir = Path(output_dir) self.output_dir.mkdir(exist_ok=True) # Initialize S3 client self.s3_client = boto3.client('s3') def _get_time_prefixes(self, start_time, end_time): """ Get timestamp prefix list to reduce scan range Args: start_time: Start timestamp end_time: End timestamp Returns: list: Timestamp prefix list, e.g. ['176005', '176006'] """ start_str = str(start_time) end_str = str(end_time) # Find first different character position min_len = min(len(start_str), len(end_str)) diff_pos = min_len for i in range(min_len): if start_str[i] != end_str[i]: diff_pos = i break # If timestamp lengths differ, take prefix of the shorter one if len(start_str) != len(end_str): prefix_len = min(len(start_str), len(end_str), diff_pos + 1) else: prefix_len = diff_pos + 1 prefixes = [] start_prefix = start_str[:prefix_len] end_prefix = end_str[:prefix_len] # Generate list of prefixes to scan if start_prefix == end_prefix: # Same prefix prefixes.append(start_prefix) else: # Different prefixes, need to scan all prefixes from start_prefix to end_prefix current = int(start_prefix) end = int(end_prefix) while current <= end: prefixes.append(str(current).zfill(prefix_len)) current += 1 return prefixes def list_files_in_time_range(self, start_time, end_time): """ List all files within specified time range (further optimized version using timestamp prefix grouping) Args: start_time: Start timestamp end_time: End timestamp Returns: list: List of matching files """ try: print(f"Scanning S3 path: s3://{self.bucket_name}/{self.prefix}") print(f"Time range: {start_time} - {end_time}") matching_files = [] # Get list of timestamp prefixes to scan time_prefixes = self._get_time_prefixes(start_time, end_time) print(f"Timestamp prefixes to scan: {time_prefixes}") total_scanned = 0 page_count = 0 # Scan each timestamp prefix for time_prefix in time_prefixes: print(f"\nScanning timestamp prefix: {time_prefix}") # Build full path for current prefix current_prefix = self.prefix + time_prefix # Build start and end markers for scanning start_marker = f"{time_prefix}*" end_marker = f"{time_prefix}*" print(f"Scanning path: s3://{self.bucket_name}/{current_prefix}") # Use pagination to get object list paginator = self.s3_client.get_paginator('list_objects_v2') for page in paginator.paginate( Bucket=self.bucket_name, Prefix=current_prefix, PaginationConfig={ 'MaxItems': None, 'PageSize': 1000 } ): page_count += 1 if page_count % 10 == 0: print(f"Scanned {page_count} pages...") if 'Contents' not in page: continue page_matched = 0 for obj in page['Contents']: total_scanned += 1 key = obj['Key'] filename = Path(key).name # Extract timestamp from filename if '-' in filename and filename.endswith('.log.gz'): try: timestamp_str = filename.split('-')[0] file_timestamp = int(timestamp_str) # Check if within time range if start_time <= file_timestamp <= end_time: matching_files.append({ 'key': key, 'filename': filename, 'size': obj['Size'], 'timestamp': file_timestamp }) page_matched += 1 except (ValueError, IndexError): continue # If this page has no matching files and enough has been scanned, consider early stop if page_matched == 0 and total_scanned > 10000: last_files = [f for f in page['Contents'][-3:] if '-' in Path(f['Key']).name] if last_files: try: last_timestamp = int(Path(last_files[-1]['Key']).name.split('-')[0]) if last_timestamp > end_time: print(f"Detected exceeding end time {end_time}, stopping scan early") break except (ValueError, IndexError): pass print(f"\nScan completed: Scanned {total_scanned} files in total, found {len(matching_files)} matching files") return matching_files except Exception as e: print(f"Error scanning S3 files: {e}") return [] def download_file(self, file_info): """ Download a single file Args: file_info: File information dictionary """ try: local_path = self.output_dir / file_info['filename'] print(f"Downloading: {file_info['filename']} ({file_info['size']} bytes)") # Download to temporary file temp_path = local_path.with_suffix('.tmp') self.s3_client.download_file( self.bucket_name, file_info['key'], str(temp_path) ) # If it's a gz file, decompress it if file_info['filename'].endswith('.gz'): print(f"Decompressing: {file_info['filename']}") with gzip.open(temp_path, 'rb') as f_in: with open(local_path.with_suffix('.log'), 'wb') as f_out: shutil.copyfileobj(f_in, f_out) temp_path.unlink() # Delete temporary file print(f"Decompression completed: {local_path.with_suffix('.log')}") else: # Rename to final filename temp_path.rename(local_path) print(f"Download completed: {local_path}") except Exception as e: print(f"Error downloading file {file_info['filename']}: {e}") def download_all_files(self, start_time, end_time, max_files=None): """ Download all files within specified time range Args: start_time: Start timestamp end_time: End timestamp max_files: Maximum number of files to download (for testing) """ files = self.list_files_in_time_range(start_time, end_time) if not files: print("No matching files found") return # Sort by timestamp files.sort(key=lambda x: x['timestamp']) print(f"\nStarting to download {len(files)} files:") print("=" * 50) downloaded_count = 0 for file_info in files: if max_files and downloaded_count >= max_files: print(f"Reached maximum download limit: {max_files}") break self.download_file(file_info) downloaded_count += 1 # Show progress every 10 files if downloaded_count % 10 == 0: print(f"Progress: {downloaded_count}/{len(files)}") print(f"\nDownload completed! Total {downloaded_count} files downloaded") print(f"Files saved in: {self.output_dir}") def main(): parser = argparse.ArgumentParser(description='Download S3 log files') parser.add_argument('--bucket', required=True, help='S3 bucket name') parser.add_argument('--prefix', required=True, help='S3 path prefix') parser.add_argument('--start-time', type=int, required=True, help='Start timestamp') parser.add_argument('--end-time', type=int, required=True, help='End timestamp') parser.add_argument('--output-dir', default='./downloads', help='Local download directory') parser.add_argument('--max-files', type=int, help='Maximum number of files to download (for testing)') parser.add_argument('--aws-profile', help='AWS config file name (if using config file)') parser.add_argument('--aws-access-key', help='AWS access key ID') parser.add_argument('--aws-secret-key', help='AWS secret access key') parser.add_argument('--region', default='us-east-1', help='AWS region') args = parser.parse_args() # Set AWS credentials if args.aws_access_key and args.aws_secret_key: os.environ['AWS_ACCESS_KEY_ID'] = args.aws_access_key os.environ['AWS_SECRET_ACCESS_KEY'] = args.aws_secret_key os.environ['AWS_DEFAULT_REGION'] = args.region elif args.aws_profile: os.environ['AWS_PROFILE'] = args.aws_profile # Create downloader and execute download downloader = S3LogDownloader(args.bucket, args.prefix, args.output_dir) downloader.download_all_files(args.start_time, args.end_time, args.max_files) if __name__ == "__main__": main()