259 lines
8.5 KiB
Python
259 lines
8.5 KiB
Python
import os
|
|
import re
|
|
import base64
|
|
import pdfplumber
|
|
import logging
|
|
from datetime import datetime, timedelta
|
|
from typing import List, Tuple, Optional
|
|
|
|
from google.oauth2.credentials import Credentials
|
|
from google_auth_oauthlib.flow import InstalledAppFlow
|
|
from googleapiclient.discovery import build
|
|
from google.auth.transport.requests import Request
|
|
|
|
from config import Config
|
|
|
|
# Load configuration
|
|
config = Config()
|
|
|
|
# Configure logging
|
|
logging.basicConfig(
|
|
level=getattr(logging, config.log_level),
|
|
format='%(asctime)s - %(levelname)s - %(message)s',
|
|
handlers=[
|
|
logging.FileHandler('invoice_extractor.log'),
|
|
logging.StreamHandler()
|
|
]
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Gmail API scope for read-only with attachments
|
|
SCOPES = ['https://www.googleapis.com/auth/gmail.readonly']
|
|
|
|
# Folder to save invoices
|
|
SAVE_FOLDER = config.save_folder
|
|
os.makedirs(SAVE_FOLDER, exist_ok=True)
|
|
|
|
|
|
def get_last_month_date_range():
|
|
"""Get the date range for the last month in Gmail query format."""
|
|
today = datetime.now()
|
|
|
|
# Calculate first day of current month
|
|
first_day_current = today.replace(day=1)
|
|
|
|
# Calculate first day of last month
|
|
if first_day_current.month == 1:
|
|
first_day_last_month = first_day_current.replace(year=first_day_current.year - 1, month=12)
|
|
else:
|
|
first_day_last_month = first_day_current.replace(month=first_day_current.month - 1)
|
|
|
|
# Calculate last day of last month
|
|
last_day_last_month = first_day_current - timedelta(days=1)
|
|
|
|
# Format dates for Gmail query (YYYY/MM/DD format)
|
|
start_date = first_day_last_month.strftime('%Y/%m/%d')
|
|
end_date = last_day_last_month.strftime('%Y/%m/%d')
|
|
|
|
logger.info(f"Scanning emails from {start_date} to {end_date}")
|
|
|
|
return f"after:{start_date} before:{end_date}"
|
|
|
|
|
|
def gmail_authenticate():
|
|
"""Authenticate with Gmail API using OAuth2."""
|
|
creds = None
|
|
if os.path.exists("token.json"):
|
|
creds = Credentials.from_authorized_user_file("token.json", SCOPES)
|
|
|
|
if not creds or not creds.valid:
|
|
if creds and creds.expired and creds.refresh_token:
|
|
try:
|
|
creds.refresh(Request())
|
|
logger.info("Credentials refreshed successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to refresh credentials: {e}")
|
|
creds = None
|
|
|
|
if not creds:
|
|
try:
|
|
flow = InstalledAppFlow.from_client_secrets_file("credentials.json", SCOPES)
|
|
creds = flow.run_local_server(port=0)
|
|
logger.info("New credentials obtained")
|
|
except Exception as e:
|
|
logger.error(f"Failed to obtain credentials: {e}")
|
|
raise
|
|
|
|
with open("token.json", "w") as token:
|
|
token.write(creds.to_json())
|
|
logger.info("Credentials saved to token.json")
|
|
|
|
return build("gmail", "v1", credentials=creds)
|
|
|
|
|
|
def search_messages(service, base_query: str, max_results: int = None) -> List[dict]:
|
|
"""Search for messages matching the query with date filtering."""
|
|
if max_results is None:
|
|
max_results = config.max_messages
|
|
|
|
# Add date filter for last month
|
|
date_filter = get_last_month_date_range()
|
|
full_query = f"{base_query} {date_filter}"
|
|
|
|
try:
|
|
results = service.users().messages().list(
|
|
userId="me",
|
|
q=full_query,
|
|
maxResults=max_results
|
|
).execute()
|
|
messages = results.get("messages", [])
|
|
logger.info(f"Found {len(messages)} messages matching query: {full_query}")
|
|
return messages
|
|
except Exception as e:
|
|
logger.error(f"Failed to search messages: {e}")
|
|
return []
|
|
|
|
|
|
def get_attachments(service, msg_id: str) -> List[Tuple[str, bytes]]:
|
|
"""Extract PDF attachments from a message."""
|
|
try:
|
|
message = service.users().messages().get(userId="me", id=msg_id).execute()
|
|
parts = message.get("payload", {}).get("parts", [])
|
|
attachments = []
|
|
|
|
for part in parts:
|
|
if part.get("filename") and part["filename"].lower().endswith(".pdf"):
|
|
att_id = part["body"].get("attachmentId")
|
|
if att_id:
|
|
try:
|
|
att = service.users().messages().attachments().get(
|
|
userId="me", messageId=msg_id, id=att_id
|
|
).execute()
|
|
data = base64.urlsafe_b64decode(att["data"])
|
|
attachments.append((part["filename"], data))
|
|
logger.debug(f"Extracted attachment: {part['filename']}")
|
|
except Exception as e:
|
|
logger.error(f"Failed to extract attachment {part['filename']}: {e}")
|
|
|
|
return attachments
|
|
except Exception as e:
|
|
logger.error(f"Failed to get attachments for message {msg_id}: {e}")
|
|
return []
|
|
|
|
|
|
def extract_text_from_pdf(pdf_bytes: bytes) -> str:
|
|
"""Extract text from PDF bytes."""
|
|
text = ""
|
|
temp_file = f"temp_{datetime.now().strftime('%Y%m%d_%H%M%S')}.pdf"
|
|
|
|
try:
|
|
with open(temp_file, "wb") as f:
|
|
f.write(pdf_bytes)
|
|
|
|
with pdfplumber.open(temp_file) as pdf:
|
|
for page_num, page in enumerate(pdf.pages):
|
|
page_text = page.extract_text()
|
|
if page_text:
|
|
text += page_text + "\n"
|
|
logger.debug(f"Extracted text from page {page_num + 1}")
|
|
|
|
return text
|
|
except Exception as e:
|
|
logger.error(f"Failed to extract text from PDF: {e}")
|
|
return ""
|
|
finally:
|
|
if os.path.exists(temp_file):
|
|
os.remove(temp_file)
|
|
|
|
|
|
def contains_vat_number(text: str) -> bool:
|
|
"""Check if text contains any of the configured VAT number patterns."""
|
|
if not text:
|
|
return False
|
|
|
|
text_upper = text.upper()
|
|
vat_patterns = config.get_vat_patterns()
|
|
|
|
for pattern in vat_patterns:
|
|
if re.search(pattern, text_upper):
|
|
logger.info(f"Found VAT number match with pattern: {pattern}")
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def save_invoice(filename: str, pdf_bytes: bytes) -> bool:
|
|
"""Save invoice PDF to the designated folder."""
|
|
try:
|
|
# Create filename with timestamp to avoid conflicts
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
name, ext = os.path.splitext(filename)
|
|
safe_filename = f"{name}_{timestamp}{ext}"
|
|
save_path = os.path.join(SAVE_FOLDER, safe_filename)
|
|
|
|
with open(save_path, "wb") as f:
|
|
f.write(pdf_bytes)
|
|
|
|
logger.info(f"✅ Saved invoice: {safe_filename}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Failed to save invoice {filename}: {e}")
|
|
return False
|
|
|
|
|
|
def main():
|
|
"""Main function to process Gmail messages and extract invoices."""
|
|
logger.info("Starting Gmail Invoice Extractor")
|
|
logger.info(f"Configuration: {config.gmail_query}, max_messages: {config.max_messages}")
|
|
|
|
try:
|
|
# Authenticate with Gmail
|
|
service = gmail_authenticate()
|
|
|
|
# Search for messages with PDF attachments from last month
|
|
messages = search_messages(service, config.gmail_query)
|
|
|
|
if not messages:
|
|
logger.info("No messages found with PDF attachments from last month")
|
|
return
|
|
|
|
processed_count = 0
|
|
saved_count = 0
|
|
|
|
for msg in messages:
|
|
try:
|
|
msg_id = msg["id"]
|
|
attachments = get_attachments(service, msg_id)
|
|
|
|
for filename, pdf_bytes in attachments:
|
|
processed_count += 1
|
|
logger.info(f"Processing attachment: {filename}")
|
|
|
|
# Extract text from PDF
|
|
text = extract_text_from_pdf(pdf_bytes)
|
|
|
|
if not text:
|
|
logger.warning(f"No text extracted from {filename}")
|
|
continue
|
|
|
|
# Check for VAT number
|
|
if contains_vat_number(text):
|
|
if save_invoice(filename, pdf_bytes):
|
|
saved_count += 1
|
|
else:
|
|
logger.info(f"❌ Skipped (no VAT match): {filename}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing message {msg.get('id', 'unknown')}: {e}")
|
|
continue
|
|
|
|
logger.info(f"Processing complete. Processed: {processed_count}, Saved: {saved_count}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Fatal error in main: {e}")
|
|
raise
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|