1
"""
2
Define import patches that are needed to fix top level entropy calls in the standard library:
3
4
1. random: Give it poison seed. Install overlays to crash if called. Check that the seed state has
5
   not advanced by asserting that the value gotten out of the rng after the snapshot restored
6
   matches the value that would be generated at poison seed state. Then remove overlays.
7
8
2. tempfile: Use a deterministic sequence of names when used at top level.
9
10
3. multiprocessing: Just give it poision entropy, the value is ordinarily used for IPC which can't
11
   happen because we have no processes.
12
13
Package import context is defined in entropy_import_context_packages.py
14
"""
15
 
16
from contextlib import contextmanager
17
from functools import wraps
18
 
19
try:
20
    from . import (
21
        entropy_import_context_packages,  # noqa: F401 Imported for side effects
22
    )
23
except ImportError:
24
    # When the python_process_pth_files compatibility flag is enabled, package entropy import
25
    # context is provided by workers-runtime-sdk via .pth files instead of being bundled here.
26
    pass
27
from .allow_entropy import (
28
    allow_bad_entropy_calls,
29
    get_bad_entropy_flag,
30
    raise_unless_entropy_allowed,
31
)
32
from .import_patch_manager import (
33
    block_calls,
34
    register_after_snapshot,
35
    register_before_first_request,
36
    register_exec_patch,
37
)
38
 
39
# Exported for snapshot backwards compatibility
40
__all__ = ["get_bad_entropy_flag"]
41
 
42
# Builtin modules:
43
# random, tempfile, and multiprocessing
44
 
45
RANDOM_STATE = None
46
 
47
 
48
@register_exec_patch("random")
49
@contextmanager
50
def random_exec(random):
51
    """Importing random calls getentropy() 10 times it seems"""
52
    with allow_bad_entropy_calls(10):
53
        yield
54
 
55
    # Block calls to functions that use the bad random seed we produced from the
56
    # ten getentropy() calls. Instantiating Random with a given seed is fine,
57
    # instantiating it without a seed will call getentropy() and fail.
58
    # Instantiating SystemRandom is fine, calling its methods will call
59
    # getentropy() and fail.
60
    global RANDOM_STATE
61
    RANDOM_STATE = random.getstate()
62
    block_calls(random, allowlist=("Random", "SystemRandom"))
63
 
64
 
65
@register_after_snapshot("random")
66
def random_after_snapshot(random):
67
    # Check that random seed hasn't been advanced somehow while executing top level scope
68
    r1 = random.random()
69
    random.setstate(RANDOM_STATE)
70
    r2 = random.random()
71
    if r1 != r2:
72
        raise RuntimeError("random seed in bad state")
73
 
74
 
75
@register_before_first_request("random")
76
def random_before_first_request(random):
77
    random.seed()
78
 
79
 
80
orig_Random_seed = None
81
 
82
 
83
@register_exec_patch("_random")
84
@contextmanager
85
def _random_exec(_random):
86
    yield
87
    global orig_Random_seed
88
    orig_Random_seed = _random.Random.seed
89
 
90
    @wraps(orig_Random_seed)
91
    def patched_seed(self, val):
92
        """
93
        Random.seed calls _PyOs_URandom which will fatally fail in top level.
94
        Prevent this by raising a RuntimeError instead.
95
        """
96
        if val is None:
97
            raise_unless_entropy_allowed()
98
        return orig_Random_seed(self, val)
99
 
100
    _random.Random.seed = patched_seed
101
 
102
 
103
@register_before_first_request("_random")
104
def _random_before_first_request(_random):
105
    _random.Random.seed = orig_Random_seed
106
 
107
 
108
class DeterministicRandomNameSequence:
109
    characters = "abcdefghijklmnopqrstuvwxyz0123456789_"
110
 
111
    def __init__(self):
112
        self.index = 0
113
 
114
    def __iter__(self):
115
        return self
116
 
117
    def index_to_chars(self):
118
        base = len(self.characters)
119
        idx = self.index
120
        s = []
121
        for _ in range(8):
122
            s.append(self.characters[idx % base])
123
            idx //= base
124
        return "".join(s)
125
 
126
    def __next__(self):
127
        self.index += 1
128
        return self.index_to_chars()
129
 
130
 
131
@register_exec_patch("tempfile")
132
@contextmanager
133
def tempfile_context(module):
134
    yield
135
    module._orig_RandomNameSequence = module._RandomNameSequence
136
    module._RandomNameSequence = DeterministicRandomNameSequence
137
 
138
 
139
@register_before_first_request("tempfile")
140
def tempfile_restore_random_name_sequence(tempfile):
141
    tempfile._RandomNameSequence = tempfile._orig_RandomNameSequence
142
    del tempfile._orig_RandomNameSequence
143
 
144
 
145
@register_exec_patch("multiprocessing.process")
146
@contextmanager
147
def multiprocessing_process_context(module):
148
    # multiprocessing.process calls os.urandom() on import. It's harmless b/c multiprocessing is
149
    # useless.
150
    with allow_bad_entropy_calls(1):
151
        yield