getCodeCLI.py

import os
import imaplib
import email
from email.message import EmailMessage
import smtplib
import argparse
import time
from dotenv import load_dotenv
import openai
from openai import OpenAI
import re

# ─── Load settings ───────────────────────────────────────────────────────────────
load_dotenv()
EMAIL_USERNAME = os.getenv('EMAIL_USERNAME')
EMAIL_PASSWORD = os.getenv('EMAIL_PASSWORD')
FORWARD_TO = os.getenv('FORWARD_TO', '').split(',')
SEARCH_SUBJECT = os.getenv('SEARCH_CRITERIA')
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')

IMAP_SERVER = 'imap.gmail.com'
IMAP_PORT = 993
SMTP_SERVER = 'smtp.gmail.com'
SMTP_PORT = 587

# Initialize OpenAI client
if OPENAI_API_KEY:
    client = OpenAI(api_key=OPENAI_API_KEY)
else:
    client = None
    print("⚠️  Warning: OPENAI_API_KEY not found in environment variables")


def test_login():
    print("→ Connecting to", IMAP_SERVER)
    m = imaplib.IMAP4_SSL(IMAP_SERVER, IMAP_PORT)
    m.debug = 4
    try:
        resp = m.login(EMAIL_USERNAME, EMAIL_PASSWORD)
        print("→ LOGIN OK:", resp)
    except imaplib.IMAP4.error as e:
        print("→ LOGIN FAILED:", repr(e))
    finally:
        m.logout()


# ─── Functions ───────────────────────────────────────────────────────────────────

def extract_code_with_openai(email_content):
    """Extract code from email content using OpenAI API."""
    if not client:
        print("❌ OpenAI client not initialized. Please set OPENAI_API_KEY.")
        return None

    try:
        prompt = f"""
        Please extract all code snippets from the following email content. 
        Return only the code blocks, properly formatted with appropriate language identifiers.
        If there are multiple code blocks, separate them clearly.
        If no code is found, return "NO_CODE_FOUND".

        Email content:
        {email_content}
        """

        response = client.chat.completions.create(
            model="gpt-3.5-turbo",  # You can change to "gpt-4" if you have access
            messages=[
                {"role": "system",
                 "content": "You are a code extraction assistant. Extract and format code snippets from text."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=2000,
            temperature=0.1
        )

        extracted_code = response.choices[0].message.content.strip()

        if extracted_code == "NO_CODE_FOUND":
            return None

        return extracted_code

    except Exception as e:
        print(f"❌ Error calling OpenAI API: {e}")
        return None


def extract_code_regex_fallback(email_content):
    """Fallback method to extract code using regex patterns."""
    # Common code block patterns
    patterns = [
        r'```[\w]*\n(.*?)\n```',  # Markdown code blocks
        r'`([^`\n]+)`',  # Inline code
        r'<code>(.*?)</code>',  # HTML code tags
        r'^\s{4,}(.+)$',  # Indented code (4+ spaces)
    ]

    code_blocks = []

    for pattern in patterns:
        matches = re.findall(pattern, email_content, re.DOTALL | re.MULTILINE)
        for match in matches:
            if match.strip() and len(match.strip()) > 10:  # Filter out short matches
                code_blocks.append(match.strip())

    return code_blocks if code_blocks else None


def print_code_to_terminal(code_content, source="OpenAI"):
    """Print extracted code to terminal with nice formatting."""
    print("\n" + "=" * 60)
    print(f"🔍 CODE EXTRACTED ({source})")
    print("=" * 60)
    print(code_content)
    print("=" * 60 + "\n")


def check_email(search_criteria=None):
    """Connect to Gmail IMAP, find the latest email matching the criteria."""
    try:
        mail = imaplib.IMAP4_SSL(IMAP_SERVER, IMAP_PORT)
        mail.login(EMAIL_USERNAME, EMAIL_PASSWORD)
        mail.select('INBOX')

        # Use provided search criteria or default to ALL
        if search_criteria:
            typ, data = mail.search(None, search_criteria)
            print(f"🔍 Searching with criteria: {search_criteria}")
        else:
            typ, data = mail.search(None, 'ALL')
            print("🔍 Getting latest email from inbox")

        if typ != 'OK':
            print(f"[Error] IMAP SEARCH returned {typ}")
            return None

        uids = data[0].split()
        if not uids:
            if search_criteria:
                print(f"No emails found matching criteria: {search_criteria}")
            else:
                print("No emails found in inbox")
            return None

        # Get the latest matching email (last UID)
        latest_uid = uids[-1]

        # Fetch the full RFC822 payload of the latest email
        typ, msg_data = mail.fetch(latest_uid, '(RFC822)')
        if typ != 'OK':
            print(f"[Error] IMAP FETCH returned {typ}")
            return None

        return email.message_from_bytes(msg_data[0][1])

    except Exception as e:
        print(f"[Error] Checking email: {e}")
        return None

    finally:
        try:
            mail.logout()
        except:
            pass


def process_email_for_code(msg):
    """Extract and process code from email message."""
    try:
        # Extract plain-text body
        body = ""
        if msg.is_multipart():
            for part in msg.walk():
                if part.get_content_type() == 'text/plain' and not part.get('Content-Disposition'):
                    charset = part.get_content_charset() or 'utf-8'
                    body = part.get_payload(decode=True).decode(charset, errors='replace')
                    break
        else:
            charset = msg.get_content_charset() or 'utf-8'
            body = msg.get_payload(decode=True).decode(charset, errors='replace')

        if not body.strip():
            print("📧 Email body is empty")
            return

        print(f"📧 Processing email: {msg['Subject']}")
        print(f"📤 From: {msg['From']}")

        # Try to extract code using OpenAI first
        extracted_code = extract_code_with_openai(body)

        if extracted_code:
            print_code_to_terminal(extracted_code, "OpenAI")
        else:
            print("🤖 No code found with OpenAI, trying regex fallback...")

            # Fallback to regex extraction
            code_blocks = extract_code_regex_fallback(body)
            if code_blocks:
                print_code_to_terminal("\n\n".join(code_blocks), "Regex Fallback")
            else:
                print("❌ No code found in email")
                print("\n📄 Email content preview:")
                print("-" * 40)
                print(body[:500] + "..." if len(body) > 500 else body)
                print("-" * 40)

    except Exception as e:
        print(f"[Error] Processing email for code: {e}")


def forward_email(msg):
    """Forwards the plain-text body of `msg` via Gmail SMTP."""
    try:
        # extract plain-text body
        body = ""
        if msg.is_multipart():
            for part in msg.walk():
                if part.get_content_type() == 'text/plain' and not part.get('Content-Disposition'):
                    charset = part.get_content_charset() or 'utf-8'
                    body = part.get_payload(decode=True).decode(charset, errors='replace')
                    break
        else:
            charset = msg.get_content_charset() or 'utf-8'
            body = msg.get_payload(decode=True).decode(charset, errors='replace')

        # compose forward
        fwd = EmailMessage()
        fwd['From'] = EMAIL_USERNAME
        fwd['To'] = ', '.join(FORWARD_TO)
        fwd['Subject'] = f"FWD: {msg['Subject']}"
        fwd.set_content(body)

        # send via Gmail SMTP
        with smtplib.SMTP(SMTP_SERVER, SMTP_PORT) as smtp:
            smtp.ehlo()
            smtp.starttls()
            smtp.login(EMAIL_USERNAME, EMAIL_PASSWORD)
            smtp.send_message(fwd)

        print("✅ Email forwarded successfully.")
    except Exception as e:
        print(f"[Error] Forwarding email: {e}")


def on_check(search_criteria=None, extract_code_only=False):
    """Check email and either extract code or forward (or both)."""
    msg = check_email(search_criteria)
    if msg:
        if extract_code_only:
            process_email_for_code(msg)
        else:
            print(f"📧 Found email: {msg['Subject']}")
            print(f"📤 From: {msg['From']}")

            # Extract and print code
            process_email_for_code(msg)

            # Also forward if FORWARD_TO is configured
            if FORWARD_TO and FORWARD_TO[0]:
                forward_email(msg)
    else:
        print("ℹ️  No matching email found.")


def main():
    parser = argparse.ArgumentParser(description='Gmail Code Extractor - Extract code from emails using OpenAI')
    parser.add_argument('--test-login', action='store_true', help='Test IMAP login credentials')
    parser.add_argument('--subject', type=str, help='Search for emails with specific subject')
    parser.add_argument('--from', type=str, dest='sender', help='Search for emails from specific sender')
    parser.add_argument('--unseen', action='store_true', help='Only search unread emails')
    parser.add_argument('--watch', type=int, metavar='SECONDS',
                        help='Watch for new emails and extract code (check every N seconds)')
    parser.add_argument('--code-only', action='store_true',
                        help='Only extract and print code, do not forward emails')

    args = parser.parse_args()

    # Build search criteria based on arguments
    search_criteria = None
    search_parts = []

    if args.subject:
        search_parts.append(f'SUBJECT "{args.subject}"')
    if args.sender:
        search_parts.append(f'FROM "{args.sender}"')
    if args.unseen:
        search_parts.append('UNSEEN')

    if search_parts:
        if len(search_parts) == 1:
            search_criteria = search_parts[0]
        else:
            search_criteria = f'({" ".join(search_parts)})'

    if args.test_login:
        test_login()
    elif args.watch:
        if search_criteria:
            print(f'👀 Watching for emails matching: {search_criteria}')
        else:
            print('👀 Watching for latest emails...')
        print(f'⏰ Checking every {args.watch} seconds')
        if not args.code_only and FORWARD_TO and FORWARD_TO[0]:
            print(f'📤 Forward to: {", ".join(FORWARD_TO)}')
        print('🤖 Will extract code using OpenAI API')
        print('Press Ctrl+C to stop\n')

        try:
            while True:
                on_check(search_criteria, extract_code_only=args.code_only)
                time.sleep(args.watch)
        except KeyboardInterrupt:
            print('\n🛑 Stopped watching for emails.')
    else:
        # Default behavior: check once
        if search_criteria:
            print(f'🔍 Searching for emails matching: {search_criteria}')
        else:
            print('🔍 Getting latest email from inbox...')
        on_check(search_criteria, extract_code_only=args.code_only)


if __name__ == '__main__':
    main()