mutt_ldap.py

   1 #!/usr/bin/env python2
   2 #
   3 # Copyright (C) 2008-2013  W. Trevor King
   4 # Copyright (C) 2012-2013  Wade Berrier
   5 # Copyright (C) 2012       Niels de Vos
   6 #
   7 # This program is free software: you can redistribute it and/or modify
   8 # it under the terms of the GNU General Public License as published by
   9 # the Free Software Foundation, either version 3 of the License, or
  10 # (at your option) any later version.
  11 #
  12 # This program is distributed in the hope that it will be useful,
  13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 # GNU General Public License for more details.
  16 #
  17 # You should have received a copy of the GNU General Public License
  18 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
  19
  20 "LDAP address searches for Mutt"
  21
  22 import codecs as _codecs
  23 import ConfigParser as _configparser
  24 import hashlib as _hashlib
  25 import json as _json
  26 import locale as _locale
  27 import logging as _logging
  28 import os.path as _os_path
  29 import pickle as _pickle
  30 import sys as _sys
  31 import time as _time
  32
  33 import ldap as _ldap
  34 import ldap.sasl as _ldap_sasl
  35
  36
  37 __version__ = '0.1'
  38
  39
  40 LOG = _logging.getLogger('mutt-ldap')
  41 LOG.addHandler(_logging.StreamHandler())
  42 LOG.setLevel(_logging.ERROR)
  43
  44
  45 class Config (_configparser.SafeConfigParser):
  46     def load(self):
  47         read_configfiles = self.read(_os_path.expanduser('~/.mutt-ldap.rc'))
  48         self._setup_defaults()
  49         LOG.info(u'loaded configuration from {0}'.format(read_configfiles))
  50
  51     def get_connection_class(self):
  52         if self.getboolean('cache', 'enable'):
  53             return CachedLDAPConnection
  54         else:
  55             return LDAPConnection
  56
  57     def _setup_defaults(self):
  58         "Setup dynamic default values"
  59         self._setup_encoding_defaults()
  60         self._setup_cache_defaults()
  61
  62     def _setup_encoding_defaults(self):
  63         default_encoding = _locale.getpreferredencoding(do_setlocale=True)
  64         for key in ['output-encoding', 'argv-encoding']:
  65             self.set(
  66                 'system', key,
  67                 self.get('system', key, raw=True) or default_encoding)
  68
  69         # HACK: convert sys.std{out,err} to Unicode (not needed in Python 3)
  70         output_encoding = self.get('system', 'output-encoding')
  71         _sys.stdout = _codecs.getwriter(output_encoding)(_sys.stdout)
  72         _sys.stderr = _codecs.getwriter(output_encoding)(_sys.stderr)
  73
  74         # HACK: convert sys.argv to Unicode (not needed in Python 3)
  75         argv_encoding = self.get('system', 'argv-encoding')
  76         _sys.argv = [unicode(arg, argv_encoding) for arg in _sys.argv]
  77
  78     def _setup_cache_defaults(self):
  79         if not self.get('cache', 'fields'):
  80             # setup a reasonable default
  81             fields = ['mail', 'cn', 'displayName']  # used by format_entry()
  82             optional_column = self.get('results', 'optional-column')
  83             if optional_column:
  84                 fields.append(optional_column)
  85             self.set('cache', 'fields', ' '.join(fields))
  86
  87
  88 CONFIG = Config()
  89 CONFIG.add_section('connection')
  90 CONFIG.set('connection', 'server', 'domaincontroller.yourdomain.com')
  91 CONFIG.set('connection', 'port', '389')  # set to 636 for default over SSL
  92 CONFIG.set('connection', 'ssl', 'no')
  93 CONFIG.set('connection', 'starttls', 'no')
  94 CONFIG.set('connection', 'basedn', 'ou=x co.,dc=example,dc=net')
  95 CONFIG.add_section('auth')
  96 CONFIG.set('auth', 'user', '')
  97 CONFIG.set('auth', 'password', '')
  98 CONFIG.set('auth', 'gssapi', 'no')
  99 CONFIG.add_section('query')
 100 CONFIG.set('query', 'filter', '') # only match entries according to this filter
 101 CONFIG.set('query', 'search-fields', 'cn displayName uid mail') # fields to wildcard search
 102 CONFIG.add_section('results')
 103 CONFIG.set('results', 'optional-column', '') # mutt can display one optional column
 104 CONFIG.add_section('cache')
 105 CONFIG.set('cache', 'enable', 'yes') # enable caching by default
 106 CONFIG.set('cache', 'path', '~/.mutt-ldap.cache') # cache results here
 107 CONFIG.set('cache', 'fields', '')  # fields to cache (if empty, setup in the main block)
 108 CONFIG.set('cache', 'longevity-days', '14') # TODO: cache results for 14 days by default
 109 CONFIG.add_section('system')
 110 # HACK: Python 2.x support, see http://bugs.python.org/issue13329#msg147475
 111 CONFIG.set('system', 'output-encoding', '')  # match .muttrc's $charset
 112 # HACK: Python 2.x support, see http://bugs.python.org/issue2128
 113 CONFIG.set('system', 'argv-encoding', '')
 114
 115
 116 class LDAPConnection (object):
 117     """Wrap an LDAP connection supporting the 'with' statement
 118
 119     See PEP 343 for details.
 120     """
 121     def __init__(self, config=None):
 122         if config is None:
 123             config = CONFIG
 124         self.config = config
 125         self.connection = None
 126
 127     def __enter__(self):
 128         self.connect()
 129         return self
 130
 131     def __exit__(self, type, value, traceback):
 132         self.unbind()
 133
 134     def connect(self):
 135         if self.connection is not None:
 136             raise RuntimeError('already connected to the LDAP server')
 137         protocol = 'ldap'
 138         if self.config.getboolean('connection', 'ssl'):
 139             protocol = 'ldaps'
 140         url = '{0}://{1}:{2}'.format(
 141             protocol,
 142             self.config.get('connection', 'server'),
 143             self.config.get('connection', 'port'))
 144         LOG.info(u'connect to LDAP server at {0}'.format(url))
 145         self.connection = _ldap.initialize(url)
 146         if (self.config.getboolean('connection', 'starttls') and
 147                 protocol == 'ldap'):
 148             self.connection.start_tls_s()
 149         if self.config.getboolean('auth', 'gssapi'):
 150             sasl = _ldap_sasl.gssapi()
 151             self.connection.sasl_interactive_bind_s('', sasl)
 152         else:
 153             self.connection.bind(
 154                 self.config.get('auth', 'user'),
 155                 self.config.get('auth', 'password'),
 156                 _ldap.AUTH_SIMPLE)
 157
 158     def unbind(self):
 159         if self.connection is None:
 160             raise RuntimeError('not connected to an LDAP server')
 161         LOG.info(u'unbind from LDAP server')
 162         self.connection.unbind()
 163         self.connection = None
 164
 165     def search(self, query):
 166         if self.connection is None:
 167             raise RuntimeError('connect to the LDAP server before searching')
 168         post = u''
 169         if query:
 170             post = u'*'
 171         fields = self.config.get('query', 'search-fields').split()
 172         filterstr = u'(|{0})'.format(
 173             u' '.join([u'({0}=*{1}{2})'.format(field, query, post) for
 174                        field in fields]))
 175         query_filter = self.config.get('query', 'filter')
 176         if query_filter:
 177             filterstr = u'(&({0}){1})'.format(query_filter, filterstr)
 178         LOG.info(u'search for {0}'.format(filterstr))
 179         msg_id = self.connection.search(
 180             self.config.get('connection', 'basedn'),
 181             _ldap.SCOPE_SUBTREE,
 182             filterstr.encode('utf-8'))
 183         res_type = None
 184         while res_type != _ldap.RES_SEARCH_RESULT:
 185             try:
 186                 res_type, res_data = self.connection.result(
 187                     msg_id, all=False, timeout=0)
 188             except _ldap.ADMINLIMIT_EXCEEDED as e:
 189                 LOG.warn(u'could not handle query results: {0}'.format(e))
 190                 break
 191             if res_data:
 192                 # use `yield from res_data` in Python >= 3.3, see PEP 380
 193                 for entry in res_data:
 194                     yield entry
 195
 196
 197 class CachedLDAPConnection (LDAPConnection):
 198     _cache_version = '{0}.0'.format(__version__)
 199
 200     def connect(self):
 201         # delay LDAP connection until we actually need it
 202         self._load_cache()
 203
 204     def unbind(self):
 205         if self.connection:
 206             super(CachedLDAPConnection, self).unbind()
 207         if self._cache:
 208             self._save_cache()
 209
 210     def search(self, query):
 211         cache_hit, entries = self._cache_lookup(query=query)
 212         if cache_hit:
 213             LOG.info(u'return cached entries for {0}'.format(query))
 214             # use `yield from res_data` in Python >= 3.3, see PEP 380
 215             for entry in entries:
 216                 yield entry
 217         else:
 218             if self.connection is None:
 219                 super(CachedLDAPConnection, self).connect()
 220             entries = []
 221             keys = self.config.get('cache', 'fields').split()
 222             for entry in super(CachedLDAPConnection, self).search(query=query):
 223                 cn,data = entry
 224                 # use dict comprehensions in Python >= 2.7, see PEP 274
 225                 cached_data = dict(
 226                     [(key, data[key]) for key in keys if key in data])
 227                 entries.append((cn, cached_data))
 228                 yield entry
 229             self._cache_store(query=query, entries=entries)
 230
 231     def _load_cache(self):
 232         path = _os_path.expanduser(self.config.get('cache', 'path'))
 233         LOG.info(u'load cache from {0}'.format(path))
 234         self._cache = {}
 235         try:
 236             data = _json.load(open(path, 'rb'))
 237         except IOError as e:  # probably "No such file"
 238             LOG.warn(u'error reading cache: {0}'.format(e))
 239         except (ValueError, KeyError) as e:  # probably a corrupt cache file
 240             LOG.warn(u'error parsing cache: {0}'.format(e))
 241         else:
 242             version = data.get('version', None)
 243             if version == self._cache_version:
 244                 self._cache = data.get('queries', {})
 245             else:
 246                 LOG.debug(u'drop outdated local cache {0} != {1}'.format(
 247                         version, self._cache_version))
 248         self._cull_cache()
 249
 250     def _save_cache(self):
 251         path = _os_path.expanduser(self.config.get('cache', 'path'))
 252         LOG.info(u'save cache to {0}'.format(path))
 253         data = {
 254             'queries': self._cache,
 255             'version': self._cache_version,
 256             }
 257         with open(path, 'wb') as f:
 258             _json.dump(data, f, indent=2, separators=(',', ': '))
 259             f.write('\n'.encode('utf-8'))
 260
 261     def _cache_store(self, query, entries):
 262         self._cache[self._cache_key(query=query)] = {
 263             'entries': entries,
 264             'time': _time.time(),
 265             }
 266
 267     def _cache_lookup(self, query):
 268         data = self._cache.get(self._cache_key(query=query), None)
 269         if data is None:
 270             return (False, data)
 271         return (True, data['entries'])
 272
 273     def _cache_key(self, query):
 274         return str((self._config_id(), query))
 275
 276     def _config_id(self):
 277         """Return a unique ID representing the current configuration
 278         """
 279         config_string = _pickle.dumps(self.config)
 280         return _hashlib.sha1(config_string).hexdigest()
 281
 282     def _cull_cache(self):
 283         cull_days = self.config.getint('cache', 'longevity-days')
 284         day_seconds = 24*60*60
 285         expire = _time.time() - cull_days * day_seconds
 286         for key in list(self._cache.keys()):  # cull the cache
 287             if self._cache[key]['time'] < expire:
 288                 LOG.debug('cull entry from cache: {0}'.format(key))
 289                 self._cache.pop(key)
 290
 291
 292 def _decode_query_data(obj):
 293     if isinstance(obj, unicode):  # e.g. cached JSON data
 294         return obj
 295     return unicode(obj, 'utf-8')
 296
 297 def format_columns(address, data):
 298     yield _decode_query_data(address)
 299     yield _decode_query_data(data.get('displayName', data['cn'])[-1])
 300     optional_column = CONFIG.get('results', 'optional-column')
 301     if optional_column in data:
 302         yield _decode_query_data(data[optional_column][-1])
 303
 304 def format_entry(entry):
 305     cn,data = entry
 306     if 'mail' in data:
 307         for m in data['mail']:
 308             # http://www.mutt.org/doc/manual/manual-4.html#ss4.5
 309             # Describes the format mutt expects: address\tname
 310             yield u'\t'.join(format_columns(m, data))
 311
 312
 313 if __name__ == '__main__':
 314     CONFIG.load()
 315
 316     if len(_sys.argv) < 2:
 317         LOG.error(u'{0}: no search string given'.format(_sys.argv[0]))
 318         _sys.exit(1)
 319
 320     query = u' '.join(_sys.argv[1:])
 321
 322     connection_class = CONFIG.get_connection_class()
 323     addresses = []
 324     with connection_class() as connection:
 325         entries = connection.search(query=query)
 326         for entry in entries:
 327             addresses.extend(format_entry(entry))
 328     print(u'{0} addresses found:'.format(len(addresses)))
 329     print(u'\n'.join(addresses))