# Copyright (C) 2017 Johnny Vestergaard <jkv@unixcluster.dk>
#
# This program 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.

# This program 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 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/>.

import base64
import logging
import binascii

from heralding.capabilities.handlerbase import HandlerBase

logger = logging.getLogger(__name__)

CRLF = '\r\n'


class Imap(HandlerBase):

  def __init__(self, options, loop):
    super().__init__(options, loop)
    self.max_tries = int(self.options['protocol_specific_data']['max_attempts'])
    self.banner = self.options['protocol_specific_data']['banner']

    self.available_commands = [
        'authenticate', 'capability', 'login', 'logout', 'noop'
    ]
    self.available_mechanisms = ['plain']

  async def execute_capability(self, reader, writer, session):
    await self._handle_session(session, reader, writer)

  async def _handle_session(
      self,
      session,
      reader,
      writer,
  ):
    await self.send_message(writer, self.banner)

    state = "Not Authenticated"
    while state != 'Logout' and session.connected:
      # An exception is raised inside await reader.readline() in case of
      # sudden connection reset.
      raw_msg = await reader.readline()
      if not raw_msg:
        break

      raw_msg_str = str(raw_msg, 'utf-8')

      cmd_msg = raw_msg_str.rstrip().split(' ', 2)
      if len(cmd_msg) == 0:
        continue
      elif len(cmd_msg) == 1:
        await self.send_message(writer, "* BAD invalid command")
        continue
      elif len(cmd_msg) == 2:
        tag = cmd_msg[0]
        cmd = cmd_msg[1]
        args = ''
      else:
        tag = cmd_msg[0]
        cmd = cmd_msg[1]
        args = cmd_msg[2]

      cmd = cmd.lower()
      if cmd not in self.available_commands:
        await self.send_message(writer, tag + " BAD invalid command")
      else:
        func_to_call = getattr(self, 'cmd_{0}'.format(cmd), None)
        if func_to_call:
          return_value = await func_to_call(session, reader, writer, tag, args)
          state = return_value
        else:
          await self.send_message(writer, tag + " BAD invalid command")
    session.end_session()

  async def cmd_authenticate(self, session, reader, writer, tag, args):
    mechanism = args.split()
    if len(mechanism) == 1:
      auth_mechanism = mechanism[0].lower()
    else:
      await self.send_message(writer, tag + ' BAD invalid command')
      return "Not Authenticated"

    if auth_mechanism in self.available_mechanisms:
      # the space after '+' is needed according to RFC
      await self.send_message(writer, '+ ')
      raw_msg = await reader.read(512)

      if auth_mechanism == 'plain':
        success, credentials = self.try_b64decode(raw_msg, session)
        # \x00 is a separator between authorization identity,
        # username and password. Authorization identity isn't used in
        # this auth mechanism, so we must have 2 \x00 symbols.(RFC 4616)
        if success and credentials.count('\x00') == 2:
          raw_msg_dec = str(base64.b64decode(raw_msg), 'utf-8')
          _, user, password = raw_msg_dec.split('\x00')
          session.add_auth_attempt(
              'plaintext', username=user, password=password)
          await self.send_message(writer, tag + ' NO Authentication failed')
        else:
          await self.send_message(writer, tag + ' BAD invalid command')
    else:
      await self.send_message(writer, tag + ' BAD invalid command')
    self.stop_if_too_many_attempts(session)
    return 'Not Authenticated'

  async def cmd_capability(self, session, reader, writer, tag, args):
    await self.send_message(writer, '* CAPABILITY IMAP4rev1 AUTH=PLAIN')
    await self.send_message(writer, tag + ' OK CAPABILITY completed')
    return 'Not Authenticated'

  async def cmd_login(self, session, reader, writer, tag, args):
    if args:
      user_cred = args.split(' ', 1)
    else:
      await self.send_message(writer, tag + ' BAD invalid command')
      return 'Not Authenticated'

    # Delete first and last quote,
    # because login and password can be sent as quoted strings
    if len(user_cred) == 1:
      user = self.strip_quotes(user_cred[0])
      password = ''
    else:
      user = self.strip_quotes(user_cred[0])
      password = self.strip_quotes(user_cred[1])

    session.add_auth_attempt('plaintext', username=user, password=password)
    await self.send_message(writer, tag + ' NO Authentication failed')
    self.stop_if_too_many_attempts(session)
    return 'Not Authenticated'

  async def cmd_logout(self, session, reader, writer, tag, args):
    await self.send_message(writer, '* BYE IMAP4rev1 Server logging out')
    await self.send_message(writer, tag + ' OK LOGOUT completed')
    return 'Logout'

  async def cmd_noop(self, session, reader, writer, tag, args):
    await self.send_message(writer, tag + ' OK NOOP completed')
    return 'Not Authenticated'

  def stop_if_too_many_attempts(self, session):
    if self.max_tries < session.get_number_of_login_attempts():
      session.end_session()

  @staticmethod
  async def send_message(writer, msg):
    message_bytes = bytes(msg + CRLF, 'utf-8')
    writer.write(message_bytes)
    await writer.drain()

  @staticmethod
  def try_b64decode(b64_str, session):
    try:
      result = base64.b64decode(b64_str)
      return True, str(result, 'utf-8')
    except binascii.Error:
      logger.warning('Error decoding base64: {0} ({1})'.format(
          binascii.hexlify(b64_str), session.id))
      return False, ''

  @staticmethod
  def strip_quotes(quoted_str):
    if quoted_str.startswith('\"') and quoted_str.endswith('\"'):
      nonquoted_str = quoted_str[1:-1]
    else:
      nonquoted_str = quoted_str
    return nonquoted_str
