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 Add :file:`mutt-ldap.py` to your ``PATH`` and add the following line
  23 to your :file:`.muttrc`::
  24
  25   set query_command = "mutt-ldap.py '%s'"
  26
  27 Search for addresses with `^t`, optionally after typing part of the
  28 name.  Configure your connection by creating :file:`~/.mutt-ldap.rc`
  29 contaning something like::
  30
  31   [connection]
  32   server = myserver.example.net
  33   basedn = ou=people,dc=example,dc=net
  34
  35 See the `CONFIG` options for other available settings.
  36 """
  37
  38 import ConfigParser as _configparser
  39 import hashlib as _hashlib
  40 import os.path as _os_path
  41 import pickle as _pickle
  42
  43 import ldap as _ldap
  44 import ldap.sasl as _ldap_sasl
  45
  46
  47 CONFIG = _configparser.SafeConfigParser()
  48 CONFIG.add_section('connection')
  49 CONFIG.set('connection', 'server', 'domaincontroller.yourdomain.com')
  50 CONFIG.set('connection', 'port', '389')  # set to 636 for default over SSL
  51 CONFIG.set('connection', 'ssl', 'no')
  52 CONFIG.set('connection', 'starttls', 'no')
  53 CONFIG.set('connection', 'basedn', 'ou=x co.,dc=example,dc=net')
  54 CONFIG.add_section('auth')
  55 CONFIG.set('auth', 'user', '')
  56 CONFIG.set('auth', 'password', '')
  57 CONFIG.set('auth', 'gssapi', 'no')
  58 CONFIG.add_section('query')
  59 CONFIG.set('query', 'filter', '') # only match entries according to this filter
  60 CONFIG.set('query', 'search-fields', 'cn displayName uid mail') # fields to wildcard search
  61 CONFIG.add_section('results')
  62 CONFIG.set('results', 'optional-column', '') # mutt can display one optional column
  63 CONFIG.add_section('cache')
  64 CONFIG.set('cache', 'enable', 'yes') # enable caching by default
  65 CONFIG.set('cache', 'path', '~/.mutt-ldap.cache') # cache results here
  66 CONFIG.set('cache', 'fields', '')  # fields to cache (if empty, setup in the main block)
  67 #CONFIG.set('cache', 'longevity-days', '14') # TODO: cache results for 14 days by default
  68 CONFIG.add_section('system')
  69 # HACK: Python 2.x support, see http://bugs.python.org/issue2128
  70 CONFIG.set('system', 'argv-encoding', 'utf-8')
  71
  72 CONFIG.read(_os_path.expanduser('~/.mutt-ldap.rc'))
  73
  74
  75 class LDAPConnection (object):
  76     """Wrap an LDAP connection supporting the 'with' statement
  77
  78     See PEP 343 for details.
  79     """
  80     def __init__(self, config=None):
  81         if config is None:
  82             config = CONFIG
  83         self.config = config
  84         self.connection = None
  85
  86     def __enter__(self):
  87         self.connect()
  88         return self
  89
  90     def __exit__(self, type, value, traceback):
  91         self.unbind()
  92
  93     def connect(self):
  94         if self.connection is not None:
  95             raise RuntimeError('already connected to the LDAP server')
  96         protocol = 'ldap'
  97         if self.config.getboolean('connection', 'ssl'):
  98             protocol = 'ldaps'
  99         url = '{0}://{1}:{2}'.format(
 100             protocol,
 101             self.config.get('connection', 'server'),
 102             self.config.get('connection', 'port'))
 103         self.connection = _ldap.initialize(url)
 104         if (self.config.getboolean('connection', 'starttls') and
 105                 protocol == 'ldap'):
 106             self.connection.start_tls_s()
 107         if self.config.getboolean('auth', 'gssapi'):
 108             sasl = _ldap_sasl.gssapi()
 109             self.connection.sasl_interactive_bind_s('', sasl)
 110         else:
 111             self.connection.bind(
 112                 self.config.get('auth', 'user'),
 113                 self.config.get('auth', 'password'),
 114                 _ldap.AUTH_SIMPLE)
 115
 116     def unbind(self):
 117         if self.connection is None:
 118             raise RuntimeError('not connected to an LDAP server')
 119         self.connection.unbind()
 120         self.connection = None
 121
 122     def search(self, query):
 123         if self.connection is None:
 124             raise RuntimeError('connect to the LDAP server before searching')
 125         post = u''
 126         if query:
 127             post = u'*'
 128         fields = self.config.get('query', 'search-fields').split()
 129         filterstr = u'(|{0})'.format(
 130             u' '.join([u'({0}=*{1}{2})'.format(field, query, post) for
 131                        field in fields]))
 132         query_filter = self.config.get('query', 'filter')
 133         if query_filter:
 134             filterstr = u'(&({0}){1})'.format(query_filter, filterstr)
 135         msg_id = self.connection.search(
 136             self.config.get('connection', 'basedn'),
 137             _ldap.SCOPE_SUBTREE,
 138             filterstr.encode('utf-8'))
 139         res_type = None
 140         while res_type != _ldap.RES_SEARCH_RESULT:
 141             try:
 142                 res_type, res_data = self.connection.result(
 143                     msg_id, all=False, timeout=0)
 144             except _ldap.ADMINLIMIT_EXCEEDED:
 145                 #print "Partial results"
 146                 break
 147             if res_data:
 148                 # use `yield from res_data` in Python >= 3.3, see PEP 380
 149                 for entry in res_data:
 150                     yield entry
 151
 152
 153 class CachedLDAPConnection (LDAPConnection):
 154     def connect(self):
 155         self._load_cache()
 156         super(CachedLDAPConnection, self).connect()
 157
 158     def unbind(self):
 159         super(CachedLDAPConnection, self).unbind()
 160         self._save_cache()
 161
 162     def search(self, query):
 163         cache_hit, entries = self._cache_lookup(query=query)
 164         if cache_hit:
 165             # use `yield from res_data` in Python >= 3.3, see PEP 380
 166             for entry in entries:
 167                 yield entry
 168         else:
 169             entries = []
 170             keys = self.config.get('cache', 'fields').split()
 171             for entry in super(CachedLDAPConnection, self).search(query=query):
 172                 cn,data = entry
 173                 # use dict comprehensions in Python >= 2.7, see PEP 274
 174                 cached_data = dict(
 175                     [(key, data[key]) for key in keys if key in data])
 176                 entries.append((cn, cached_data))
 177                 yield entry
 178             self._cache_store(query=query, entries=entries)
 179
 180     def _load_cache(self):
 181         path = _os_path.expanduser(self.config.get('cache', 'path'))
 182         try:
 183             self._cache = _pickle.load(open(path, 'rb'))
 184         except IOError:  # probably "No such file"
 185             self._cache = {}
 186         except (ValueError, KeyError):  # probably a corrupt cache file
 187             self._cache = {}
 188
 189     def _save_cache(self):
 190         path = _os_path.expanduser(self.config.get('cache', 'path'))
 191         _pickle.dump(self._cache, open(path, 'wb'))
 192
 193     def _cache_store(self, query, entries):
 194         self._cache[self._cache_key(query=query)] = entries
 195
 196     def _cache_lookup(self, query):
 197         entries = self._cache.get(self._cache_key(query=query), None)
 198         if entries is None:
 199             return (False, entries)
 200         return (True, entries)
 201
 202     def _cache_key(self, query):
 203         return (self._config_id(), query)
 204
 205     def _config_id(self):
 206         """Return a unique ID representing the current configuration
 207         """
 208         config_string = _pickle.dumps(self.config)
 209         return _hashlib.sha1(config_string).hexdigest()
 210
 211
 212 def format_columns(address, data):
 213     yield unicode(address, 'utf-8')
 214     yield unicode(data.get('displayName', data['cn'])[-1], 'utf-8')
 215     optional_column = CONFIG.get('results', 'optional-column')
 216     if optional_column in data:
 217         yield unicode(data[optional_column][-1], 'utf-8')
 218
 219 def format_entry(entry):
 220     cn,data = entry
 221     if 'mail' in data:
 222         for m in data['mail']:
 223             # http://www.mutt.org/doc/manual/manual-4.html#ss4.5
 224             # Describes the format mutt expects: address\tname
 225             yield u'\t'.join(format_columns(m, data))
 226
 227
 228 if __name__ == '__main__':
 229     import sys
 230
 231     # HACK: convert sys.argv to Unicode (not needed in Python 3)
 232     argv_encoding = CONFIG.get('system', 'argv-encoding')
 233     sys.argv = [unicode(arg, argv_encoding) for arg in sys.argv]
 234
 235     if len(sys.argv) < 2:
 236         sys.stderr.write(u'{0}: no search string given\n'.format(sys.argv[0]))
 237         sys.exit(1)
 238
 239     query = u' '.join(sys.argv[1:])
 240
 241     if CONFIG.getboolean('cache', 'enable'):
 242         connection_class = CachedLDAPConnection
 243         if not CONFIG.get('cache', 'fields'):
 244             # setup a reasonable default
 245             fields = ['mail', 'cn', 'displayName']  # used by format_entry()
 246             optional_column = CONFIG.get('results', 'optional-column')
 247             if optional_column:
 248                 fields.append(optional_column)
 249             CONFIG.set('cache', 'fields', ' '.join(fields))
 250     else:
 251         connection_class = LDAPConnection
 252
 253     addresses = []
 254     with connection_class() as connection:
 255         entries = connection.search(query=query)
 256         for entry in entries:
 257             addresses.extend(format_entry(entry))
 258     print(u'{0} addresses found:'.format(len(addresses)))
 259     print(u'\n'.join(addresses))