#!/usr/bin/env python3

import argparse # Parsing arguments
import sys # Mainly for exit
import json # JSON parsing
import logging # Logging
import msal # Microsoft Authentication Library
import requests # HTTP Requests
from pprint import pprint # Pretty printing
import imaplib # Imap
import base64 # Base64

# Microsoft graph endpoint
graph_endpoint = 'https://graph.microsoft.com/v1.0{0}'

# Set up argument parsing
parser = argparse.ArgumentParser(description="Checks IMAP status of a mailbox")
parser.add_argument("-t", "--tenant-id", dest="tenant", help="Azure AD tenant id", required=True)
parser.add_argument("-c", "--client-id", dest="clientid", help="Azure application client ID", required=True)
parser.add_argument("-s", "--secret", dest="secret", help="Azure application secret")
parser.add_argument("-e", "--email", dest="email", help="Mailbox to log into", required=True)
parser.add_argument("-E", "--server", dest="server", help="Specify mail server", required=True)
parser.add_argument("-d", "--debug", dest="debug", action="store_true", help="Enable debugging output")
parser.add_argument("-T", "--thumbprint", dest="thumbprint", help="Certificate thumbprint if using cert-based authentication")
parser.add_argument("-p", "--private-key", dest="pkey", help="Location of private key")

args = parser.parse_args()

# Set up debug
if args.debug:
    log_level = getattr(logging, "DEBUG")
else:
    log_level = getattr(logging, "INFO")

# Set up logging
logging.basicConfig(level=log_level)
logger = logging.getLogger("check_message_count")

# Verify that if thumbprint is provided, private key is provided and vice versa
if (args.thumbprint is None and args.pkey is not None) or (args.thumbprint is not None and args.pkey is None):
    print("[UNKNOWN] - -T/--thumbprint and -p/--private-key must both be provided when using certificate-based authentication.")
    sys.exit(3)

# Verify that not trying to use cert and secret at the same time
if any([args.thumbprint, args.pkey]) and args.secret is not None:
    print("[UNKNOWN] - Either -s/--secret must be provided or -T/--thumbprint AND -p/--private-key")
    sys.exit(3)
    
# Verify that we aren't just using nothing
if not any([args.thumbprint, args.pkey, args.secret]):
    print("[UNKNOWN] - Either -s/--secret must be provided or -T/--thumbprint AND -p/--private-key")
    sys.exit(3)
    
if args.thumbprint is not None:
    with open(args.pkey, 'r', encoding='utf-8') as f:
        creds = {"thumbprint": args.thumbprint, "private_key": f.read()}
else:
    creds = args.secret
    

# Get MSAL app
app = msal.ConfidentialClientApplication(
    args.clientid, authority=f"https://login.microsoftonline.com/{args.tenant}",
    client_credential=creds
)

# Declare scope
scope = ["https://graph.microsoft.com/.default"]

# Try to get app.
result = app.acquire_token_silent(scope, account=None)

# If no app exists, create one.

if not result:
    logging.debug("No suitable token exists in cache, acquiring new token from AAD.")
    result = app.acquire_token_for_client(scopes=scope)

if not result or 'access_token' not in result:
    print("CRITICAL - Unable to acquire access token.")
    pprint(result)
    sys.exit(3)

# Function definitions
def make_api_get_call(url, token, parameters=None):
    headers = {
        "Authorization": f'Bearer {token}',
        "Accept": "application/json"
    }
    logger.debug(headers)
    logger.debug(url)
    logger.debug(parameters)
    return requests.get(url, headers=headers, params=parameters, timeout=10)

def get_messages(access_token, user, top="30"):
    # Format string
    get_messages_url = graph_endpoint.format(f'/users/{user}/mailfolders/inbox/messages')

    # Query parameters
    query_parameters = {
        "$top": top,
        "$select": "receivedDateTime,subject,from",
        "$orderby": "receivedDateTime DESC"
    }
    logger.debug(get_messages_url)
    logger.debug(query_parameters)
    logger.debug(access_token)
    return make_api_get_call(get_messages_url, access_token, parameters=query_parameters)


# Get list of emails
email_results = json.loads(get_messages(result['access_token'], args.email).text)['value']
val = len(email_results)

if (val > 0):
    print(f"CRITICAL - Found {val} emails")
    sys.exit(2)
elif (val == 0):
    print("OK - No emails found.")
    sys.exit(0)
else:
    print("UNKNOWN - Something weird happened.")
    sys.exit(3)