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

# Created by Stuart Colville on 2008-02-02
# Muffin Research Labs. http://muffinresearch.co.uk/
# Updated by Graham Dutton on 2008-11-23
#
# Copyright (c) 2008, Stuart J Colville
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the Muffin Research Labs nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY Stuart J Colville ``AS IS'' AND ANY
# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL Stuart J Colville BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import os, sys, commands, re

def listkeychains():
  """ Returns a dictionary of all of the keychains found on the system """
  k = commands.getoutput("security list-keychains")
  rx = re.compile(r'".*?/([\w]*)\.keychain"', re.I | re.M)
  keychains = {}
  for match in rx.finditer(k):
    keychains[match.group(1)]=match.group().strip('"')
  return keychains

def createkeychain(newkeychain, password=None):
  if not password:
     from getpass import getpass
     password = getpass('Password:')

  newkeychain = newkeychain.find('.keychain') > -1  and newkeychain or '%s.keychain' % newkeychain
  k = commands.getstatusoutput("security create-keychain -p %s %s" % (password, newkeychain,))
  if self.DEBUG:
    print k
  if k[0]:
    return False, 'Create creation failed'
  if k[0] == False:
    return keychain(newkeychain), 'Keychain created successfully'

def checkkeychainname(keychain, list):
   """ Rationalises keychain strings as to whether they have .keychain or not
   and looks them up in the dictionary of keychains created at
   instantiation. Returns a string if successful and False if keychain is
   not available"""

   ext = '.keychain'
   if not keychain.find(ext): keychain += ext

   if keychain in list or keychain[0] == '/':
      return keychain
   return False

class Keychain:
  DEBUG=False

  def __init__(self, keychainname):
    """ Keychain.py  is a simple class allowing access to keychain data and
    settings. Keychain.py can also setup new keychains as required. As the
    keychain is only available on MaxOSX the module will raise ImportError
    if import is attempted on anything other than Mac OSX """

    if sys.platform == 'darwin':
      self.keychains = listkeychains()
    else:
      raise ImportError('Keychain is only available on Mac OSX')
    self.keychain = checkkeychainname(keychainname, self.keychains)
    if not self.keychain:
      raise ValueError('Keychain not found')

  def getinternetpassword(self, item):
    """ Returns account + password pair from specified keychain item """
    return self._getpassword(servername=item, type='internet')

  def getgenericpassword(self, item):
    """ Returns account + password pair from specified keychain item """
    return self._getpassword(accountname=item, type='generic')

  def _getpassword(self, accountname=None, servername=None, type='generic'):
    args="security find-%s-password -g " % type
    if accountname:
        args += "-a %s " % accountname
        match = 'acct'
    if servername:
        args += "-s %s " % servername
        match = 'srvr'
    args += "%s " % self.keychain

    k = commands.getstatusoutput(args)
    if self.DEBUG:
      print k
    if k[0]:
      return False, 'The specified item could not be found', k
    else:
      rx1 = re.compile(r'"'+match+'"<blob>="(.*?)"', re.S)
      rx2 = re.compile(r'password: "(.*?)"', re.S)
      account = rx1.search(k[1])
      password = rx2.search(k[1])
      if account and password:
        return {match:account.group(1),"password":password.group(1)}
      else:
        return False

  def setgenericpassword(self, account, password, servicename=None):
    """ Create and store a generic account and  password in the given keychain """
    account = account and '-a %s' % (account,) or ''
    password = password and '-p %s' % (password,) or ''
    servicename = servicename and '-s %s' % (servicename,) or ''
    k = commands.getstatusoutput("security add-generic-password %s %s %s %s" % (account, password, servicename, self.keychain))
    if self.DEBUG:
      print k
    if k[0]:
      return False, 'The specified password could not be added to %s' % self.keychain
    if k[0] == False:
      return True, 'Password added to %s successfully' % self.keychain

  def lockkeychain(self):
    k = commands.getstatusoutput("security lock-keychain %s" % (self.keychain,))
    if self.DEBUG:
      print k
    if k[0]:
      return False, 'Keychain: %s could not be locked' % self.keychain
    if k[0] == False:
      return True, 'Keychain: %s locked successfully' % self.keychain


  def unlockkeychain(self, password=None):
    if not password:
       from getpass import getpass
       password = getpass('Password:')
    k = commands.getstatusoutput("security unlock-keychain -p %s %s" % (password, self.keychain,))
    if self.DEBUG:
      print k
    if k[0]:
      return False, 'Keychain could not be unlocked'
    if k[0] == False:
      return True, 'Keychain unlocked successfully'

  def setkeychain(self, lock=True, timeout=0):
    """ Allows setting the keychain configuration. If lock is True the
    keychain will be locked on sleep. If the timeout is set to anything other
    than 0 the keychain will be set to lock after timeout seconds of
    inactivity """

    lock = lock and '-l' or ''
    timeout = timeout and '-u -t %s' % (timeout,) or ''
    k = commands.getstatusoutput("security keychain-settings %s %s %s" % (lock, timeout, self.keychain,))
    if self.DEBUG:
      print k
    if k[0]:
      return False, 'Keychain settings failed'
    if k[0] == False:
      return True, 'Keychain updated successfully'

  def showkeychaininfo(self):
    """Returns a dictionary containing the keychain settings"""

    k = commands.getstatusoutput("security show-keychain-info %s" % (self.keychain,))
    if self.DEBUG:
      print k
    if k[0]:
      return False, 'Keychain could not be found'
    if k[0] == False:
      result = {}
      result['keychain'] = self.keychain
      if k[1].find('lock-on-sleep') > -1:
        result['lock-on-sleep'] = True
      if k[1].find('no-timeout') > -1:
        result['timeout'] = 0
      else:
        rx = re.compile(r'timeout=(\d+)s', re.S)
        match = rx.search(k[1])
        result['timeout'] = match.group(1)
      return result