#! /usr/bin/env python
# -*- coding: utf-8 -*-

# imapt -- Simple command-line tool to test IMAP server
# By: Benjamin RENARD <brenard@easter-eggs.com>
#
# Copyright (C) 2013 Easter-eggs
# http://git.zionetrix.net/mailt
#
# This file is part of MailT.
#
# MailT is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published
# by the Free Software Foundation, either version 3 of the License,
# or (at your option) any later version.
#
# MailT is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from optparse import OptionParser

import os
import imaplib

import logging
import sys

import re
import getpass

parser = OptionParser()


parser.add_option('-H',
		  '--host',
		  action="store",
		  type="string",
		  dest="host",
		  help="The IMAP host",
		  default="127.0.0.1")

parser.add_option('-p',
		  '--port',
		  action="store",
		  type="string",
		  dest="port",
		  help="The IMAP port",
		  default=None)

parser.add_option('-S',
		  '--ssl',
		  action="store_true",
		  dest="ssl",
		  help="IMAP Over SSL (IMAPS)",
		  default=False)

parser.add_option('-U',
		  '--user',
		  action="store",
		  type="string",
		  dest="user",
		  help="User login",
		  default=None)

parser.add_option('-P',
		  '--password',
		  action="store",
		  type="string",
		  dest="password",
		  help="User password",
		  default=None)

parser.add_option('--md5',
		  action="store_true",
		  dest="md5",
		  help="Use CRAM-MD5 password for authentication",
		  default=False)

parser.add_option('-m',
		  '--mailbox',
		  action="store",
		  type="string",
		  dest="mailbox",
		  help="The mailbox directory to select",
		  default='INBOX')

parser.add_option('-v',
		  '--verbose',
		  action="store_true",
		  dest="verbose")

parser.add_option('-l',
		  '--list',
		  action="store_true",
		  dest="list",
		  help="List mailboxes",
		  default=False)

parser.add_option('-q',
		  '--quota',
		  action="store_true",
		  dest="quota",
		  help="Display quota informations",
		  default=False)

parser.add_option('--expunge',
		  action="store",
		  type="string",
		  dest="expunge",
		  help="Expunge mailbox specify (ALL to expunge all mailboxes)",
		  default=None)

(options, args) = parser.parse_args()

if options.verbose:
	loglevel=logging.DEBUG
else:
	loglevel=logging.INFO

logging.basicConfig(level=loglevel,format='%(asctime)s - %(levelname)s - %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p')

if not options.user:
	logging.fatal('You must provide the user login')
	sys.exit(1)

if not options.password:
	options.password=getpass.getpass()

if options.port is None:
	if options.ssl:
		options.port=993
	else:
		options.port=143

if options.ssl:
	logging.debug('Establish SSL connexion to server')
	try:
		server = imaplib.IMAP4_SSL(options.host,options.port)
		logging.debug('Connected')
	except Exception, e:
		logging.critical("IMAP connection error : %s" % e)
		sys.exit(2)
else:
	logging.debug("Establish connexion to server")
	try:
		server = imaplib.IMAP4(options.host,options.port)
		logging.debug('Connected')
	except Exception, e:
		logging.critical("IMAP connection error : %s" % e)
		sys.exit(2)

if options.user:
	try:
		if options.md5:
			logging.debug('Login into server as %s (with CRAM-MD5 password)' % options.user)
			server.login_cram_md5(options.user,options.password)
		else:
			logging.debug('Login into server as %s' % options.user)
			server.login(options.user,options.password)
		logging.debug('Logged')
	except Exception as e:
		logging.critical("IMAP Auth error : %s" % e)
		sys.exit(3)

def expunge(server,mailbox):
	logging.info('Expunge %s mailbox' % mailbox)
	try:
		(st,c) = server.select(mailbox)
		if st!='OK':
			logging.error('Error selecting %s mailbox : %s' % (mailbox,st))
			return
		(st,ex) = server.expunge()
		if st=='OK':
			if len(ex)==1 and ex[0] is None:
				logging.info('Mailbox %s expunged, no mail removed' % mailbox)
			else:
				logging.info('Mailbox %s expunged : %s mail(s) removed' % (mailbox,len(ex)))
		else:
			logging.info('Error expunging %s mailbox, return "%s"' % (mailbox,st))
	except Exception as e:
		logging.error('Error expunging %s mailbox : %s' % (mailbox,e))

def format_quota(num):
	num=int(num)
	for x in ['KB', 'MB', 'GB', 'TB']:
		if num < 1024.0:
			return "%3.1f %s" % (num, x)
		num /= 1024.0

try:
	if options.list:
		logging.info('List mailbox')
		(status,l) = server.list()
		for d in l:
			logging.info('   "%s"' % d.split('"')[3])
	elif options.expunge:
		if options.expunge=='ALL':
			logging.debug("Listing mailboxes")
			(status,l) = server.list()
			for d in l:
				dirname=d.split('"')[3]
				expunge(server,dirname)
		else:
			expunge(server,options.expunge)
	elif options.quota:
		logging.debug("Get quota usage")
		ret=server.getquotaroot(options.mailbox)
		logging.debug("IMAP server return : %s" % str(ret))
		# Return example on Dovecot
		# ('OK', [['"INBOX" "User quota"'], ['"User quota" (STORAGE 0 1024)']])
		if isinstance(ret, (list, tuple)):
			for quota in ret[1][1]:
				if quota is None:
					logging.info('No quota defined')
					continue
				parsequota=re.match('^"([^"]*)" \(STORAGE ([0-9]*) ([0-9]*)\)$',quota)
				if parsequota:
					logging.info("Quota %s : %s / %s" % (parsequota.group(1),format_quota(parsequota.group(2)),format_quota(parsequota.group(3))))
				else:
					logging.error("Error parsing quota string from IMAP server : '%s'" % quota)
		else:
			logging.error("Invalid result getting quotas: %s" % ret)
	else:
		logging.debug("Select mailbox")
		ret = server.select(options.mailbox)
		if len(ret)>0 and ret[0]=='OK':
			logging.info("Mailbox %s content %s message(s)" % (options.mailbox,ret[1][0]))
		else:
			logging.error("Selecting return incorrected : %s" % ret)
		server.close()
except Exception as e:
	logging.critical('Error selecting mailbox %s : %s' % (options.mailbox,e))
finally:
	server.logout()