1#!/usr/bin/env python2
2# -*- coding: utf-8 -*-
3
4from __future__ import division, absolute_import, print_function, unicode_literals
5import io
6import ldap
7import ldap.modlist
8import ldif
9import sys
10import hashlib
11from argparse import ArgumentParser
12from ldif import LDIFWriter
13
14
15def _safe_encode(data):
16    if isinstance(data, unicode):
17        return data.encode('utf-8')
18    return data
19
20
21class BytesLDIFRecordList(ldif.LDIFRecordList):
22    """Simple encoding wrapper for LDIFRecordList that converts keys to UTF-8"""
23    def _next_key_and_value(self):
24        # we do not descend from object, so we cannot use super()
25        k, v = ldif.LDIFRecordList._next_key_and_value(self)
26        return k.encode('utf-8'), v
27
28
29def ldap_connect(address, binddn, password):
30    try:
31        conn = ldap.initialize("ldap://%s" % address, bytes_mode=True)
32    except TypeError:
33        conn = ldap.initialize("ldap://%s" % address)
34
35    conn.set_option(ldap.OPT_REFERRALS, 0)
36    conn.simple_bind_s(binddn, password)
37
38    return conn
39
40
41def ldap_create_or_fail(conn, dn, modlist):
42    """
43    create an object dn with the attributes from modlist dict
44    """
45    try:
46        conn.add_s(dn, ldap.modlist.addModlist(modlist))
47    except ldap.ALREADY_EXISTS:
48        print("Object '%s' already exists." % dn, file=sys.stderr)
49        sys.exit(1)
50
51
52def action_clean(conn, basedn):
53    """
54    Clean up all objects in our subtrees (ldap-variant of "rm -rf").
55
56    find the dns ob all objects below our bases and remove them ordered from
57    longest to shortest dn, so we remove parent objects later
58    """
59
60    for subtree_dn in [b"ou=backup,%s" % basedn, b"ou=restore,%s" % basedn]:
61        try:
62            for dn in sorted(
63                map(
64                    lambda x: x[0],
65                    conn.search_s(subtree_dn, ldap.SCOPE_SUBTREE, attrsonly=1),
66                ),
67                key=len,
68                reverse=True,
69            ):
70                conn.delete_s(dn)
71        except ldap.NO_SUCH_OBJECT:
72            # if the top object doesn't exist, there's nothing to remove
73            pass
74
75
76def action_populate(conn, basedn):
77    """Populate our backup data"""
78    ldap_create_or_fail(
79        conn,
80        b"ou=backup,%s" % basedn,
81        {b"objectClass": [b"organizationalUnit"], b"ou": [b"restore"]},
82    )
83
84    ldap_create_or_fail(
85        conn,
86        b"cn=No JPEG,ou=backup,%s" % basedn,
87        {
88            b"objectClass": [b"inetOrgPerson", b"posixAccount", b"shadowAccount"],
89            b"uid": [b"njpeg"],
90            b"sn": [b"JPEG"],
91            b"givenName": [b"No"],
92            b"cn": [b"No JPEG"],
93            b"displayName": [b"No JPEG"],
94            b"uidNumber": [b"1000"],
95            b"gidNumber": [b"1000"],
96            b"loginShell": [b"/bin/bash"],
97            b"homeDirectory": [b"/home/njpeg"],
98        },
99    )
100
101    ldap_create_or_fail(
102        conn,
103        b"cn=Small JPEG,ou=backup,%s" % basedn,
104        {
105            b"objectClass": [b"inetOrgPerson", b"posixAccount", b"shadowAccount"],
106            b"uid": [b"sjpeg"],
107            b"sn": [b"JPEG"],
108            b"givenName": [b"Small"],
109            b"cn": [b"Small JPEG"],
110            b"displayName": [b"Small JPEG"],
111            b"uidNumber": [b"1001"],
112            b"gidNumber": [b"1000"],
113            b"loginShell": [b"/bin/bash"],
114            b"homeDirectory": [b"/home/sjpeg"],
115            b"jpegPhoto": open("image-small.jpg", "rb").read(),
116        },
117    )
118
119    ldap_create_or_fail(
120        conn,
121        b"cn=Medium JPEG,ou=backup,%s" % basedn,
122        {
123            b"objectClass": [b"inetOrgPerson", b"posixAccount", b"shadowAccount"],
124            b"uid": [b"mjpeg"],
125            b"sn": [b"JPEG"],
126            b"givenName": [b"Medium"],
127            b"cn": [b"Medium JPEG"],
128            b"displayName": [b"Medium JPEG"],
129            b"uidNumber": [b"1002"],
130            b"gidNumber": [b"1000"],
131            b"loginShell": [b"/bin/bash"],
132            b"homeDirectory": [b"/home/mjpeg"],
133            b"jpegPhoto": open("image-medium.jpg", "rb").read(),
134        },
135    )
136
137    ldap_create_or_fail(
138        conn,
139        b"cn=Large JPEG,ou=backup,%s" % basedn,
140        {
141            b"objectClass": [b"inetOrgPerson", b"posixAccount", b"shadowAccount"],
142            b"uid": [b"ljpeg"],
143            b"sn": [b"JPEG"],
144            b"givenName": [b"Large"],
145            b"cn": [b"Large JPEG"],
146            b"displayName": [b"Large JPEG"],
147            b"uidNumber": [b"1003"],
148            b"gidNumber": [b"1000"],
149            b"loginShell": [b"/bin/bash"],
150            b"homeDirectory": [b"/home/ljpeg"],
151            b"jpegPhoto": open("image-large.jpg", "rb").read(),
152        },
153    )
154
155    ldap_create_or_fail(
156        conn,
157        b"o=Bareos GmbH & Co. KG,ou=backup,%s" % basedn,
158        {b"objectClass": [b"top", b"organization"], b"o": [b"Bareos GmbH & Co. KG"]},
159    )
160
161    ldap_create_or_fail(
162        conn,
163        b"ou=automount,ou=backup,%s" % basedn,
164        {b"objectClass": [b"top", b"organizationalUnit"], b"ou": [b"automount"]},
165    )
166
167    # # Objects with / in the DN are currently not supported
168    # ldap_create_or_fail(
169    #     conn,
170    #     "cn=/home,ou=automount,ou=backup,%s" % basedn,
171    #     {
172    #         "objectClass": [b"top", b"person"],
173    #         "cn": [b"/home"],
174    #         "sn": [b"Automount objects don't have a surname"],
175    #     },
176    # )
177
178    ldap_create_or_fail(
179        conn,
180        b"ou=weird-names,ou=backup,%s" % basedn,
181        {b"objectClass": [b"top", b"organizationalUnit"], b"ou": [b"weird-names"]},
182    )
183
184    for ou in [
185        b" leading-space",
186        b"#leading-hash",
187        b"space in middle",
188        b"trailing-space ",
189        b"with\nnewline",
190        b"with,comma",
191        b'with"quotes"',
192        b"with\\backslash",
193        b"with+plus",
194        b"with#hash",
195        b"with;semicolon",
196        b"with<less-than",
197        b"with=equals",
198        b"with>greater-than",
199    ]:
200        ldap_create_or_fail(
201            conn,
202            b"ou=%s,ou=weird-names,ou=backup,%s" % (ldap.dn.escape_dn_chars(ou), basedn),
203            {b"objectClass": [b"top", b"organizationalUnit"], b"ou": [ou]},
204        )
205
206    # creating the DN using the normal method wouldn't work, so we create a
207    # temporary LDIF and parse that.
208    ldif_data = io.BytesIO()
209    ldif_data.write(b"dn: ou=böses encoding,ou=weird-names,ou=backup,")
210    ldif_data.write(b"%s\n" % basedn.encode('ascii'))
211    ldif_data.write(b"objectClass: top\n")
212    ldif_data.write(b"objectClass: organizationalUnit\n")
213    ldif_data.write(b"ou: böses encoding\n")
214    ldif_data.seek(0)
215
216    ldif_parser = BytesLDIFRecordList(ldif_data, max_entries=1)
217    ldif_parser.parse()
218    dn, entry = ldif_parser.all_records[0]
219    ldif_data.close()
220
221    ldap_create_or_fail(
222        conn,
223        _safe_encode(dn),
224        entry
225        )
226
227
228def abbrev_value(v):
229    """Abbreviate long values for readable LDIF output"""
230    length = len(v)
231    if length > 80:
232        digest = hashlib.sha1(v).hexdigest()
233        return "BLOB len:%d sha1:%s" % (length, digest)
234    return v
235
236
237def action_dump(conn, basedn, shorten=True, rewrite_dn=True):
238    writer = LDIFWriter(sys.stdout)
239    try:
240        for dn, attrs in conn.search_s(basedn, ldap.SCOPE_SUBTREE):
241            if rewrite_dn:
242                dn = (
243                    dn.decode("utf-8")
244                    .replace(basedn, "dc=unified,dc=base,dc=dn")
245                    .encode("utf-8")
246                )
247            if shorten:
248                attrs = {
249                    k: [abbrev_value(v) for v in vals] for k, vals in attrs.iteritems()
250                }
251            try:
252                writer.unparse(dn, attrs)
253            except UnicodeDecodeError:
254                writer.unparse(dn.decode('utf-8'), attrs)
255    except ldap.NO_SUCH_OBJECT:
256        print("No object '%s' in directory." % basedn, file=sys.stderr)
257        sys.exit(1)
258
259
260if __name__ == "__main__":
261    parser = ArgumentParser(description="Tool to create, remove LDAP test data")
262    parser.add_argument(
263        "--clean", action="store_true", help="remove data from LDAP server"
264    )
265    parser.add_argument(
266        "--populate", action="store_true", help="populate LDAP server with data"
267    )
268    parser.add_argument(
269        "--dump-backup",
270        action="store_true",
271        help="print representation of backup subtree",
272    )
273    parser.add_argument(
274        "--dump-restore",
275        action="store_true",
276        help="print representation of restore subtree",
277    )
278    parser.add_argument(
279        "--full-value",
280        action="store_true",
281        help="disable shortening of large values during dump",
282    )
283    parser.add_argument(
284        "--real-dn", action="store_true", help="disable rewriting of DN during dump"
285    )
286    parser.add_argument(
287        "--address", "--host", default="localhost", help="LDAP server address"
288    )
289    parser.add_argument(
290        "--basedn", "-b", default=b"dc=example,dc=org", help="LDAP base dn"
291    )
292    parser.add_argument(
293        "--binddn", "-D", default=b"cn=admin,dc=example,dc=org", help="LDAP bind dn"
294    )
295    parser.add_argument("--password", "-w", default=b"admin", help="LDAP password")
296
297    args = parser.parse_args()
298
299    if (
300        not args.clean
301        and not args.populate
302        and not args.dump_backup
303        and not args.dump_restore
304    ):
305        print("please select at least one action", file=sys.stderr)
306        sys.exit(1)
307
308    conn = ldap_connect(args.address, args.binddn, args.password)
309    if args.clean:
310        action_clean(conn, args.basedn)
311    if args.populate:
312        action_populate(conn, args.basedn)
313    if args.dump_backup:
314        action_dump(
315            conn,
316            b"ou=backup,%s" % args.basedn,
317            shorten=not args.full_value,
318            rewrite_dn=not args.real_dn,
319        )
320    if args.dump_restore:
321        action_dump(
322            conn,
323            b"ou=restore,%s" % args.basedn,
324            shorten=not args.full_value,
325            rewrite_dn=not args.real_dn,
326        )
327