Explorer
/tmp/restic-stage/lead-engine/gmail_integration.py
← Zurück ↓ Download
"""
Gmail API Integration
Handles reading and sending emails via Gmail API
"""

import os
import base64
from typing import List, Optional, Tuple
from email.mime.text import MIMEText
from google.auth.transport.requests import Request
from google.oauth2.service_account import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
import pickle
import googleapiclient.discovery
from googleapiclient.errors import HttpError


class GmailClient:
    """Handles Gmail API operations"""

    SCOPES = ['https://www.googleapis.com/auth/gmail.modify']

    def __init__(self, credentials_file: str = 'credentials.json',
                 token_file: str = 'token.pickle'):
        """
        Initialize Gmail client with OAuth2 credentials

        Args:
            credentials_file: Path to Google OAuth credentials.json
            token_file: Path to store OAuth token (pickle format)
        """
        self.credentials_file = credentials_file
        self.token_file = token_file
        self.service = None
        self._authenticate()

    def _authenticate(self):
        """Authenticate with Gmail API using OAuth2"""
        creds = None

        # Load existing token if available
        if os.path.exists(self.token_file):
            with open(self.token_file, 'rb') as token:
                creds = pickle.load(token)

        # If no valid credentials, get new ones
        if not creds or not creds.valid:
            if creds and creds.expired and creds.refresh_token:
                creds.refresh(Request())
            else:
                flow = InstalledAppFlow.from_client_secrets_file(
                    self.credentials_file, self.SCOPES)
                creds = flow.run_local_server(port=0)

            # Save token for next run
            with open(self.token_file, 'wb') as token:
                pickle.dump(creds, token)

        self.service = googleapiclient.discovery.build('gmail', 'v1', credentials=creds)

    def get_unread_messages(self, max_results: int = 10) -> List[dict]:
        """
        Fetch unread emails from inbox

        Args:
            max_results: Maximum number of messages to fetch

        Returns:
            List of message dictionaries with email data
        """
        try:
            results = self.service.users().messages().list(
                userId='me',
                q='is:unread',
                maxResults=max_results
            ).execute()

            messages = results.get('messages', [])
            email_data = []

            for msg in messages:
                msg_data = self.get_message_details(msg['id'])
                if msg_data:
                    email_data.append(msg_data)

            return email_data

        except HttpError as error:
            print(f'An error occurred: {error}')
            return []

    def get_message_details(self, message_id: str) -> Optional[dict]:
        """
        Get full details of a specific message

        Args:
            message_id: Gmail message ID

        Returns:
            Dictionary with email details (from, subject, body, etc.)
        """
        try:
            message = self.service.users().messages().get(
                userId='me',
                id=message_id,
                format='full'
            ).execute()

            headers = message['payload']['headers']
            subject = next((h['value'] for h in headers if h['name'] == 'Subject'), 'No Subject')
            from_email = next((h['value'] for h in headers if h['name'] == 'From'), 'Unknown')
            date = next((h['value'] for h in headers if h['name'] == 'Date'), '')

            # Extract body
            body = self._extract_body(message['payload'])

            # Parse from_email and from_name
            from_name, from_addr = self._parse_from_header(from_email)

            return {
                'message_id': message_id,
                'from_email': from_addr,
                'from_name': from_name,
                'subject': subject,
                'body': body,
                'timestamp': date,
                'raw': message
            }

        except HttpError as error:
            print(f'An error occurred: {error}')
            return None

    def _extract_body(self, payload: dict) -> str:
        """Extract body text from message payload"""
        if 'parts' in payload:
            # Multipart message
            for part in payload['parts']:
                if part['mimeType'] == 'text/plain':
                    data = part['body'].get('data', '')
                    if data:
                        return base64.urlsafe_b64decode(data).decode('utf-8')
        else:
            # Simple message
            data = payload['body'].get('data', '')
            if data:
                return base64.urlsafe_b64decode(data).decode('utf-8')

        return ''

    def _parse_from_header(self, from_header: str) -> Tuple[str, str]:
        """
        Parse From header into name and email

        Args:
            from_header: Raw From header value

        Returns:
            Tuple of (name, email)
        """
        # Format: "Name <email@example.com>" or "email@example.com"
        if '<' in from_header and '>' in from_header:
            name = from_header.split('<')[0].strip().strip('"')
            email = from_header.split('<')[1].split('>')[0].strip()
            return (name, email)
        else:
            return ('', from_header.strip())

    def send_email(self, to: str, subject: str, body: str) -> bool:
        """
        Send an email via Gmail

        Args:
            to: Recipient email address
            subject: Email subject
            body: Email body (plain text)

        Returns:
            True if sent successfully, False otherwise
        """
        try:
            message = MIMEText(body)
            message['to'] = to
            message['subject'] = subject

            raw = base64.urlsafe_b64encode(message.as_bytes()).decode()

            self.service.users().messages().send(
                userId='me',
                body={'raw': raw}
            ).execute()

            return True

        except HttpError as error:
            print(f'An error occurred: {error}')
            return False

    def mark_as_read(self, message_id: str) -> bool:
        """Mark a message as read"""
        try:
            self.service.users().messages().modify(
                userId='me',
                id=message_id,
                body={'removeLabelIds': ['UNREAD']}
            ).execute()
            return True
        except HttpError as error:
            print(f'An error occurred: {error}')
            return False

    def add_label(self, message_id: str, label_name: str) -> bool:
        """Add a label to a message"""
        try:
            # First get or create label
            labels = self.service.users().labels().list(userId='me').execute()
            label_id = None

            for label in labels.get('labels', []):
                if label['name'] == label_name:
                    label_id = label['id']
                    break

            # Create label if not found
            if not label_id:
                label_body = {
                    'name': label_name,
                    'labelListVisibility': 'labelShow',
                    'messageListVisibility': 'show'
                }
                created_label = self.service.users().labels().create(
                    userId='me',
                    body=label_body
                ).execute()
                label_id = created_label['id']

            # Add label to message
            self.service.users().messages().modify(
                userId='me',
                id=message_id,
                body={'addLabelIds': [label_id]}
            ).execute()

            return True

        except HttpError as error:
            print(f'An error occurred: {error}')
            return False