1import sys
2
3import salt.payload
4
5
6def _trim_dict_in_dict(data, max_val_size, replace_with):
7    """
8    Takes a dictionary, max_val_size and replace_with
9    and recursively loops through and replaces any values
10    that are greater than max_val_size.
11    """
12    for key in data:
13        if isinstance(data[key], dict):
14            _trim_dict_in_dict(data[key], max_val_size, replace_with)
15        else:
16            if sys.getsizeof(data[key]) > max_val_size:
17                data[key] = replace_with
18
19
20def trim_dict(
21    data,
22    max_dict_bytes,
23    percent=50.0,
24    stepper_size=10,
25    replace_with="VALUE_TRIMMED",
26    is_msgpacked=False,
27    use_bin_type=False,
28):
29    """
30    Takes a dictionary and iterates over its keys, looking for
31    large values and replacing them with a trimmed string.
32
33    If after the first pass over dictionary keys, the dictionary
34    is not sufficiently small, the stepper_size will be increased
35    and the dictionary will be rescanned. This allows for progressive
36    scanning, removing large items first and only making additional
37    passes for smaller items if necessary.
38
39    This function uses msgpack to calculate the size of the dictionary
40    in question. While this might seem like unnecessary overhead, a
41    data structure in python must be serialized in order for sys.getsizeof()
42    to accurately return the items referenced in the structure.
43
44    Ex:
45    >>> salt.utils.dicttrim.trim_dict({'a': 'b', 'c': 'x' * 10000}, 100)
46    {'a': 'b', 'c': 'VALUE_TRIMMED'}
47
48    To improve performance, it is adviseable to pass in msgpacked
49    data structures instead of raw dictionaries. If a msgpack
50    structure is passed in, it will not be unserialized unless
51    necessary.
52
53    If a msgpack is passed in, it will be repacked if necessary
54    before being returned.
55
56    :param use_bin_type: Set this to true if "is_msgpacked=True"
57                         and the msgpack data has been encoded
58                         with "use_bin_type=True". This also means
59                         that the msgpack data should be decoded with
60                         "encoding='utf-8'".
61    """
62    if is_msgpacked:
63        dict_size = sys.getsizeof(data)
64    else:
65        dict_size = sys.getsizeof(salt.payload.dumps(data))
66    if dict_size > max_dict_bytes:
67        if is_msgpacked:
68            if use_bin_type:
69                data = salt.payload.loads(data, encoding="utf-8")
70            else:
71                data = salt.payload.loads(data)
72        while True:
73            percent = float(percent)
74            max_val_size = float(max_dict_bytes * (percent / 100))
75            try:
76                for key in data:
77                    if isinstance(data[key], dict):
78                        _trim_dict_in_dict(data[key], max_val_size, replace_with)
79                    else:
80                        if sys.getsizeof(data[key]) > max_val_size:
81                            data[key] = replace_with
82                percent = percent - stepper_size
83                max_val_size = float(max_dict_bytes * (percent / 100))
84                if use_bin_type:
85                    dump_data = salt.payload.dumps(data, use_bin_type=True)
86                else:
87                    dump_data = salt.payload.dumps(data)
88                cur_dict_size = sys.getsizeof(dump_data)
89                if cur_dict_size < max_dict_bytes:
90                    if is_msgpacked:  # Repack it
91                        return dump_data
92                    else:
93                        return data
94                elif max_val_size == 0:
95                    if is_msgpacked:
96                        return dump_data
97                    else:
98                        return data
99            except ValueError:
100                pass
101        if is_msgpacked:
102            if use_bin_type:
103                return salt.payload.dumps(data, use_bin_type=True)
104            else:
105                return salt.payload.dumps(data)
106        else:
107            return data
108    else:
109        return data
110