import os
import re
import requests
import json
from datetime import datetime
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Retrieve API key and management console URL securely
api_url = os.getenv("MANAGEMENT_CONSOLE_URL") + "/web/api/v2.1/threats?limit=10"
api_key = os.getenv("API_KEY")
log_file_path = "/var/log/sentinelone.json"
custom_timestamp = ""  # Enter your preferred timestamp using format 2023-01-01T00:00:00


def get_last_timestamp(log_file_path):
    try:
        with open(log_file_path, 'r') as file:
            lines = file.readlines()
            if lines:
                last_line = lines[-1].strip()
                match = re.search(r'"createdAt":\s*"([^"]+)"', last_line)
                if match:
                    last_created_at = match.group(1)
                    last_timestamp = datetime.strptime(last_created_at, "%Y-%m-%dT%H:%M:%S.%fZ").isoformat()
                    return last_timestamp
                else:
                    return None
            else:
                return None
    except FileNotFoundError:
        return None


def get_logs(start_timestamp):
    headers = {
        'Authorization': f'ApiToken {api_key}',
        'Content-Type': 'application/json'
    }

    # Construct query parameters
    params = {}
    if start_timestamp:
        params['createdAt__gt'] = start_timestamp

    response = requests.get(api_url, headers=headers, params=params)
    if response.status_code == 200:
        return response.json()
    else:
        print(f"Failed to fetch logs: {response.status_code}")
        return None


def main():
    if not api_key or not api_url:
        print("Error: API key or Management Console URL not set. Check your .env file.")
        return

    # Get the last timestamp from the log file
    last_timestamp = get_last_timestamp(log_file_path)
    if last_timestamp:
        print(f"Last timestamp in log file: {last_timestamp}")
    else:
        print("Log file is empty or doesn't exist.")

    if custom_timestamp:
        last_timestamp_from_file = get_last_timestamp(log_file_path)
        start_timestamp = last_timestamp_from_file if last_timestamp_from_file else custom_timestamp
        print(f"Using timestamp: {start_timestamp}")
    else:
        start_timestamp = last_timestamp if last_timestamp else None
        print(f"Using last timestamp from log file: {start_timestamp}" if start_timestamp else "No last timestamp found.")

    # Query the SentinelOne API for logs since the start timestamp
    logs = get_logs(start_timestamp)

    if logs:
        with open(log_file_path, 'a') as file:
            for log in logs['data']:
                file.write(json.dumps(log))
                file.write('\n')
        print(f"Logs written to {log_file_path}")
    else:
        print("No logs fetched.")


if __name__ == "__main__":
    main()