#!/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
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("-S", "--subject", dest="subject", help="Subject to verify against", required=True)
parser.add_argument("-E", "--server", dest="server", help="Specify mail server", required=True)
parser.add_argument("-f", "--from", dest="sender", help="Specify required sender address")
parser.add_argument("-n", "--no-delete", dest="nodelete", help="Don't delete found items", action="store_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, "WARN")

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

# 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 'error' in result:
        print(f"[UNKNOWN] - {result['error_description']}")
        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)

def delete_message(access_token, user, message_id):
    # Headers for request
    headers = {
        "Authorization": f'Bearer {access_token}',
        "Accept": "application/json"
    }

    # Define message delete url
    delete_message_url = graph_endpoint.format(f'/users/{user}/messages/{message_id}')

    #Make call and get result
    result = requests.delete(delete_message_url, headers=headers, timeout=10)

    logger.debug(result)
    logger.debug(result.status_code)
    logger.debug(result.text)
    if result.status_code == 204:
        return True
    else:
        return False


# Get list of emails
email_results = json.loads(get_messages(result['access_token'], args.email).text)
logger.debug(f"\nResult:\n{pprint.pformat(email_results, indent=4)}")
if 'error' in email_results:
    print(f"[UNKNOWN] - Error encountered while obtaining email: {email_results['error']['message']}")
    sys.exit(3)
email_results = email_results['value']
emails_to_delete = 0
emails_deleted = 0

# Cycle through emails for deletion status
for email in email_results:
    if email['subject'].lower() == args.subject.lower():
        if args.sender is not None and args.sender.lower() != email['from']['emailAddress']['address'].lower():
            continue
        else:
            emails_to_delete += 1
            logger.debug(email)
            logger.debug(type(result))
            logger.debug(type(args.email))
            logger.debug(type(email))

            # Delete message
            if args.nodelete:
                continue
            if delete_message(result['access_token'], args.email, email['id']):
                emails_deleted += 1


# Status code
if (emails_to_delete == 0):
    print("[CRITICAL] - Found 0 emails")
    sys.exit(2)
elif (emails_to_delete != emails_deleted):
    print(f"[CRITICAL] - Script was unable to delete all emails. {emails_to_delete} found, {emails_deleted} deleted.")
    sys.exit(2)
elif (emails_to_delete == emails_deleted) and (emails_to_delete >= 1):
    print(f"[OK] - {emails_to_delete} found, {emails_deleted} deleted.")
    sys.exit(0) 
else:
    print("[UNKNOWN] - Something went wrong.")
    sys.exit(3)