summaryrefslogtreecommitdiff
path: root/src/leap/mail/walk.py
blob: 820b8c77b9370ee0ef77d1cf0bbde8b884e25652 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
# -*- coding: utf-8 -*-
# walk.py
# Copyright (C) 2013 LEAP
#
# 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/>.
"""
Utilities for walking along a message tree.
"""
import hashlib
import os

from leap.mail.utils import first

DEBUG = os.environ.get("BITMASK_MAIL_DEBUG")

if DEBUG:
    get_hash = lambda s: hashlib.sha256(s).hexdigest()[:10]
else:
    get_hash = lambda s: hashlib.sha256(s).hexdigest()


"""
Get interesting message parts
"""
get_parts = lambda msg: [
    {'multi': part.is_multipart(),
     'ctype': part.get_content_type(),
     'size': len(part.as_string()),
     'parts': len(part.get_payload())
        if isinstance(part.get_payload(), list)
        else 1,
     'headers': part.items(),
     'phash': get_hash(part.get_payload())
        if not part.is_multipart() else None}
    for part in msg.walk()]

"""
Utility lambda functions for getting the parts vector and the
payloads from the original message.
"""

get_parts_vector = lambda parts: (x.get('parts', 1) for x in parts)
get_payloads = lambda msg: ((x.get_payload(),
                             dict(((str.lower(k), v) for k, v in (x.items()))))
                            for x in msg.walk())

get_body_phash_simple = lambda payloads: first(
    [get_hash(payload) for payload, headers in payloads
     if "text/plain" in headers.get('content-type')])

get_body_phash_multi = lambda payloads: (first(
    [get_hash(payload) for payload, headers in payloads
     if "text/plain" in headers.get('content-type')])
    or get_body_phash_simple(payloads))

"""
On getting the raw docs, we get also some of the headers to be able to
index the content. Here we remove any mutable part, as the the filename
in the content disposition.
"""

get_raw_docs = lambda msg, parts: (
    {"type": "cnt",  # type content they'll be
     "raw": payload if not DEBUG else payload[:100],
     "phash": get_hash(payload),
     "content-disposition": first(headers.get(
         'content-disposition', '').split(';')),
     "content-type": headers.get(
         'content-type', ''),
     "content-transfer-encoding": headers.get(
         'content-transfer-type', '')}
    for payload, headers in get_payloads(msg)
    if not isinstance(payload, list))


def walk_msg_tree(parts, body_phash=None):
    """
    Take a list of interesting items of a message subparts structure,
    and return a dict of dicts almost ready to be written to the content
    documents that will be stored in Soledad.

    It walks down the subparts in the parsed message tree, and collapses
    the leaf docuents into a wrapper document until no multipart submessages
    are left. To achieve this, it iteratively calculates a wrapper vector of
    all documents in the sequence that have more than one part and have unitary
    documents to their right. To collapse a multipart, take as many
    unitary documents as parts the submessage contains, and replace the object
    in the sequence with the new wrapper document.

    :param parts: A list of dicts containing the interesting properties for
                  the message structure. Normally this has been generated by
                  doing a message walk.
    :type parts: list of dicts.
    :param body_phash: the payload hash of the body part, to be included
                       in the outer content doc for convenience.
    :type body_phash: basestring or None
    """
    # parts vector
    pv = list(get_parts_vector(parts))

    if len(parts) == 2:
        inner_headers = parts[1].get("headers", None)

    if DEBUG:
        print "parts vector: ", pv
        print

    # wrappers vector
    getwv = lambda pv: [True if pv[i] != 1 and pv[i + 1] == 1 else False
                        for i in range(len(pv) - 1)]
    wv = getwv(pv)

    # do until no wrapper document is left
    while any(wv):
        wind = wv.index(True)  # wrapper index
        nsub = pv[wind]  # number of subparts to pick
        slic = parts[wind + 1:wind + 1 + nsub]  # slice with subparts

        cwra = {
            "multi": True,
            "part_map": dict((index + 1, part)  # content wrapper
                             for index, part in enumerate(slic)),
            "headers": dict(parts[wind]['headers'])
        }

        # remove subparts and substitue wrapper
        map(lambda i: parts.remove(i), slic)
        parts[wind] = cwra

        # refresh vectors for this iteration
        pv = list(get_parts_vector(parts))
        wv = getwv(pv)

    outer = parts[0]
    outer.pop('headers')
    if not "part_map" in outer:
        # we have a multipart with 1 part only, so kind of fix it
        # although it would be prettier if I take this special case at
        # the beginning of the walk.
        pdoc = {"multi": True,
                "part_map": {1: outer}}
        pdoc["part_map"][1]["multi"] = False
        if not pdoc["part_map"][1].get("phash", None):
            pdoc["part_map"][1]["phash"] = body_phash
        pdoc["part_map"][1]["headers"] = inner_headers
    else:
        pdoc = outer
    pdoc["body"] = body_phash
    return pdoc