workspace_client.py

import os
import json
import datetime
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
from google.auth.transport.requests import Request
from googleapiclient.discovery import build
import config

class WorkspaceClient:
    def __init__(self):
        self.credentials = self._get_credentials()
        self.service = build('admin', 'reports_v1', credentials=self.credentials)

    def _get_credentials(self):
        """
        Authenticates using OAuth 2.0 desktop/installed app flow.
        - On first run, opens a browser for the admin to sign in.
        - Caches the token in token.json for future runs.
        - No service account key needed.
        """
        creds = None

        # Check for cached token
        if os.path.exists(config.TOKEN_FILE):
            creds = Credentials.from_authorized_user_file(config.TOKEN_FILE, config.SCOPES)

        # If no valid creds, do the OAuth flow
        if not creds or not creds.valid:
            if creds and creds.expired and creds.refresh_token:
                print("  Refreshing access token...")
                creds.refresh(Request())
            else:
                if not os.path.exists(config.OAUTH_CLIENT_FILE):
                    raise FileNotFoundError(
                        f"OAuth client file not found: {config.OAUTH_CLIENT_FILE}\n"
                        f"Download it from GCP Console → APIs & Services → Credentials → "
                        f"OAuth 2.0 Client IDs → Download JSON.\n"
                        f"Save it as '{config.OAUTH_CLIENT_FILE}' in this directory."
                    )
                print("  Opening browser for Google sign-in...")
                flow = InstalledAppFlow.from_client_secrets_file(
                    config.OAUTH_CLIENT_FILE, config.SCOPES
                )
                creds = flow.run_local_server(port=0)

            # Cache the token
            with open(config.TOKEN_FILE, 'w') as token:
                token.write(creds.to_json())
            print("  ✔ Authenticated successfully.")

        return creds

    def _get_start_time(self, days_back=1):
        """Returns an RFC 3339 timestamp for `days_back` days ago."""
        now = datetime.datetime.utcnow()
        start_date = now - datetime.timedelta(days=days_back)
        return start_date.isoformat() + 'Z'

    def _get_end_time(self):
        """Returns current time as RFC 3339."""
        return datetime.datetime.utcnow().isoformat() + 'Z'

    # Apps that require both startTime AND endTime
    _NEEDS_END_TIME = {"gmail"}
    # Apps with a max lookback stricter than 30 days
    _MAX_DAYS = {"gmail": 29}

    def fetch_logs(self, application_name, user_key='all', max_results=None,
                   start_time=None, days_back=1):
        """
        Fetches audit logs for a specific application.
        :param application_name: e.g., 'login', 'admin', 'token', 'drive', 'gmail'
        :param user_key: A user email to filter by, or 'all' for everyone.
        :param max_results: Maximum number of events to return.
        :param start_time: RFC 3339 string. Overrides days_back if provided.
        :param days_back: How many days of logs to fetch (default 1).
        """
        if max_results is None:
            max_results = config.MAX_RESULTS_PER_FETCH

        # Enforce per-app max lookback
        cap = self._MAX_DAYS.get(application_name)
        if cap and days_back > cap:
            days_back = cap

        if start_time is None:
            start_time = self._get_start_time(days_back)

        label = f"{application_name} logs for '{user_key}'" if user_key != 'all' else f"{application_name} logs"
        print(f"  Fetching {label} (past {days_back}d)...")

        try:
            params = dict(
                userKey=user_key,
                applicationName=application_name,
                maxResults=max_results,
                startTime=start_time,
            )
            # Gmail (and potentially others) require endTime
            if application_name in self._NEEDS_END_TIME:
                params['endTime'] = self._get_end_time()

            results = self.service.activities().list(**params).execute()

            activities = results.get('items', [])
            return activities

        except Exception as e:
            print(f"  Error fetching logs for {application_name}: {e}")
            return []

    def fetch_all_monitored_logs(self):
        """
        Fetches logs for all applications defined in config.APPLICATIONS_TO_MONITOR.
        Returns a dictionary mapping app names to their log events.
        """
        all_logs = {}
        for app in config.APPLICATIONS_TO_MONITOR:
            logs = self.fetch_logs(app)
            all_logs[app] = logs
        return all_logs

    def fetch_user_investigation_logs(self, user_email, days_back=7):
        """
        Fetches all monitored log types for a specific user.
        Defaults to 7 days for investigations (vs 1 day for routine scans).
        Returns a dict of {app_name: [events]}.
        """
        print(f"\n  Gathering investigation data for {user_email} (past {days_back} days)...")
        all_logs = {}
        total = 0
        for app in config.APPLICATIONS_TO_MONITOR:
            logs = self.fetch_logs(app, user_key=user_email, days_back=days_back)
            all_logs[app] = logs
            total += len(logs)
        print(f"  Total events collected: {total}")
        return all_logs