1from __future__ import absolute_import
2
3"""M2Crypto wrapper for OpenSSL BIO API.
4
5Copyright (c) 1999-2004 Ng Pheng Siong. All rights reserved."""
6
7import io
8import logging
9from typing import Any, AnyStr, Callable, Iterable, Optional, Union  # noqa
10
11from M2Crypto import m2, six
12
13log = logging.getLogger('BIO')
14
15
16class BIOError(ValueError):
17    pass
18
19
20m2.bio_init(BIOError)
21
22
23class BIO(object):
24    """Abstract object interface to the BIO API."""
25
26    m2_bio_free = m2.bio_free
27
28    def __init__(self, bio=None, _pyfree=0, _close_cb=None):
29        # type: (Optional[BIO], int, Optional[Callable]) -> None
30        self.bio = bio
31        self._pyfree = _pyfree
32        self._close_cb = _close_cb
33        self.closed = 0
34        self.write_closed = 0
35
36    def __del__(self):
37        if self._pyfree:
38            self.m2_bio_free(self.bio)
39
40    def _ptr(self):
41        return self.bio
42
43    # Deprecated.
44    bio_ptr = _ptr
45
46    def fileno(self):
47        # type: () -> int
48        return m2.bio_get_fd(self.bio)
49
50    def readable(self):
51        # type: () -> bool
52        return not self.closed
53
54    def read(self, size=None):
55        # type: (int) -> Union[bytes, bytearray]
56        if not self.readable():
57            raise IOError('cannot read')
58        if size is None:
59            buf = bytearray()
60            while 1:
61                data = m2.bio_read(self.bio, 4096)
62                if not data:
63                    break
64                buf += data
65            return buf
66        elif size == 0:
67            return b''
68        elif size < 0:
69            raise ValueError('read count is negative')
70        else:
71            return bytes(m2.bio_read(self.bio, size))
72
73    def readline(self, size=4096):
74        # type: (int) -> bytes
75        if not self.readable():
76            raise IOError('cannot read')
77        buf = m2.bio_gets(self.bio, size)
78        buf = '' if buf is None else buf
79        return six.ensure_binary(buf)
80
81    def readlines(self, sizehint='ignored'):
82        # type: (Union[AnyStr, int]) -> Iterable[bytes]
83        if not self.readable():
84            raise IOError('cannot read')
85        lines = []
86        while 1:
87            buf = m2.bio_gets(self.bio, 4096)
88            if buf is None:
89                break
90            lines.append(six.ensure_binary(buf))
91        return lines
92
93    def writeable(self):
94        # type: () -> bool
95        return (not self.closed) and (not self.write_closed)
96
97    def write(self, data):
98        # type: (AnyStr) -> int
99        """Write data to BIO.
100
101        :return: either data written, or [0, -1] for nothing written,
102                 -2 not implemented
103        """
104        if not self.writeable():
105            raise IOError('cannot write')
106        if isinstance(data, six.text_type):
107            data = data.encode('utf8')
108        return m2.bio_write(self.bio, data)
109
110    def write_close(self):
111        # type: () -> None
112        self.write_closed = 1
113
114    def flush(self):
115        # type: () -> None
116        """Flush the buffers.
117
118        :return: 1 for success, and 0 or -1 for failure
119        """
120        m2.bio_flush(self.bio)
121
122    def reset(self):
123        # type: () -> int
124        """Set the bio to its initial state.
125
126        :return: 1 for success, and 0 or -1 for failure
127        """
128        return m2.bio_reset(self.bio)
129
130    def close(self):
131        # type: () -> None
132        self.closed = 1
133        if self._close_cb:
134            self._close_cb()
135
136    def should_retry(self):
137        # type: () -> int
138        """
139        Can the call be attempted again, or was there an error
140        ie do_handshake
141
142        """
143        return m2.bio_should_retry(self.bio)
144
145    def should_read(self):
146        # type: () -> int
147        """Should we read more data?"""
148
149        return m2.bio_should_read(self.bio)
150
151    def should_write(self):
152        # type: () -> int
153        """Should we write more data?"""
154        return m2.bio_should_write(self.bio)
155
156    def tell(self):
157        """Return the current offset."""
158        return m2.bio_tell(self.bio)
159
160    def seek(self, off):
161        """Seek to the specified absolute offset."""
162        return m2.bio_seek(self.bio, off)
163
164    def __enter__(self):
165        return self
166
167    def __exit__(self, *args):
168        # type: (*Any) -> int
169        self.close()
170
171
172class MemoryBuffer(BIO):
173    """Object interface to BIO_s_mem.
174
175    Empirical testing suggests that this class performs less well than
176    cStringIO, because cStringIO is implemented in C, whereas this class
177    is implemented in Python. Thus, the recommended practice is to use
178    cStringIO for regular work and convert said cStringIO object to
179    a MemoryBuffer object only when necessary.
180    """
181
182    def __init__(self, data=None):
183        # type: (Optional[bytes]) -> None
184        super(MemoryBuffer, self).__init__(self)
185        if data is not None and not isinstance(data, bytes):
186            raise TypeError(
187                "data must be bytes or None, not %s" % (type(data).__name__, ))
188        self.bio = m2.bio_new(m2.bio_s_mem())
189        self._pyfree = 1
190        if data is not None:
191            m2.bio_write(self.bio, data)
192
193    def __len__(self):
194        # type: () -> int
195        return m2.bio_ctrl_pending(self.bio)
196
197    def read(self, size=0):
198        # type: (int) -> bytes
199        if not self.readable():
200            raise IOError('cannot read')
201        if size:
202            return m2.bio_read(self.bio, size)
203        else:
204            return m2.bio_read(self.bio, m2.bio_ctrl_pending(self.bio))
205
206    # Backwards-compatibility.
207    getvalue = read_all = read
208
209    def write_close(self):
210        # type: () -> None
211        super(MemoryBuffer, self).write_close()
212        m2.bio_set_mem_eof_return(self.bio, 0)
213
214    close = write_close
215
216
217class File(BIO):
218    """Object interface to BIO_s_pyfd.
219
220    This class interfaces Python to OpenSSL functions that expect BIO. For
221    general file manipulation in Python, use Python's builtin file object.
222    """
223
224    def __init__(self, pyfile, close_pyfile=1, mode='rb'):
225        # type: (Union[io.BytesIO, AnyStr], int, AnyStr) -> None
226        super(File, self).__init__(self, _pyfree=1)
227
228        if isinstance(pyfile, six.string_types):
229            pyfile = open(pyfile, mode)
230
231        # This is for downward compatibility, but I don't think, that it is
232        # good practice to have two handles for the same file. Whats about
233        # concurrent write access? Last write, last wins? Especially since Py3
234        # has its own buffer management. See:
235        #
236        #  https://docs.python.org/3.3/c-api/file.html
237        #
238        pyfile.flush()
239        self.fname = pyfile.name
240        self.pyfile = pyfile
241        # Be wary of https://github.com/openssl/openssl/pull/1925
242        # BIO_new_fd is NEVER to be used before OpenSSL 1.1.1
243        if hasattr(m2, "bio_new_pyfd"):
244            self.bio = m2.bio_new_pyfd(pyfile.fileno(), m2.bio_noclose)
245        else:
246            self.bio = m2.bio_new_pyfile(pyfile, m2.bio_noclose)
247
248        self.close_pyfile = close_pyfile
249        self.closed = False
250
251    def flush(self):
252        # type: () -> None
253        super(File, self).flush()
254        self.pyfile.flush()
255
256    def close(self):
257        # type: () -> None
258        self.flush()
259        super(File, self).close()
260        if self.close_pyfile:
261            self.pyfile.close()
262
263    def reset(self):
264        # type: () -> int
265        """Set the bio to its initial state.
266
267        :return: 0 for success, and -1 for failure
268        """
269        return super(File, self).reset()
270
271    def __del__(self):
272        if not self.closed:
273            m2.bio_free(self.bio)
274
275
276def openfile(filename, mode='rb'):
277    # type: (AnyStr, AnyStr) -> File
278    try:
279        f = open(filename, mode)
280    except IOError as ex:
281        raise BIOError(ex.args)
282
283    return File(f)
284
285
286class IOBuffer(BIO):
287    """Object interface to BIO_f_buffer.
288
289    Its principal function is to be BIO_push()'ed on top of a BIO_f_ssl, so
290    that makefile() of said underlying SSL socket works.
291    """
292
293    m2_bio_pop = m2.bio_pop
294    m2_bio_free = m2.bio_free
295
296    def __init__(self, under_bio, mode='rwb', _pyfree=1):
297        # type: (BIO, str, int) -> None
298        super(IOBuffer, self).__init__(self, _pyfree=_pyfree)
299        self.io = m2.bio_new(m2.bio_f_buffer())
300        self.bio = m2.bio_push(self.io, under_bio._ptr())
301        # This reference keeps the underlying BIO alive while we're not closed.
302        self._under_bio = under_bio
303        if 'w' in mode:
304            self.write_closed = 0
305        else:
306            self.write_closed = 1
307
308    def __del__(self):
309        # type: () -> None
310        if getattr(self, '_pyfree', 0):
311            self.m2_bio_pop(self.bio)
312        self.m2_bio_free(self.io)
313
314    def close(self):
315        # type: () -> None
316        BIO.close(self)
317
318
319class CipherStream(BIO):
320    """Object interface to BIO_f_cipher."""
321
322    SALT_LEN = m2.PKCS5_SALT_LEN
323
324    m2_bio_pop = m2.bio_pop
325    m2_bio_free = m2.bio_free
326
327    def __init__(self, obio):
328        # type: (BIO) -> None
329        super(CipherStream, self).__init__(self, _pyfree=1)
330        self.obio = obio
331        self.bio = m2.bio_new(m2.bio_f_cipher())
332        self.closed = 0
333
334    def __del__(self):
335        # type: () -> None
336        if not getattr(self, 'closed', 1):
337            self.close()
338
339    def close(self):
340        # type: () -> None
341        self.m2_bio_pop(self.bio)
342        self.m2_bio_free(self.bio)
343        self.closed = 1
344
345    def write_close(self):
346        # type: () -> None
347        self.obio.write_close()
348
349    def set_cipher(self, algo, key, iv, op):
350        # type: (str, AnyStr, AnyStr, int) -> None
351        cipher = getattr(m2, algo, None)
352        if cipher is None:
353            raise ValueError('unknown cipher', algo)
354        else:
355            if not isinstance(key, bytes):
356                key = key.encode('utf8')
357            if not isinstance(iv, bytes):
358                iv = iv.encode('utf8')
359        m2.bio_set_cipher(self.bio, cipher(), key, iv, int(op))
360        m2.bio_push(self.bio, self.obio._ptr())
361
362
363class SSLBio(BIO):
364    """Object interface to BIO_f_ssl."""
365
366    def __init__(self, _pyfree=1):
367        # type: (int) -> None
368        super(SSLBio, self).__init__(self, _pyfree=_pyfree)
369        self.bio = m2.bio_new(m2.bio_f_ssl())
370        self.closed = 0
371
372    def set_ssl(self, conn, close_flag=m2.bio_noclose):
373        ## type: (Connection, int) -> None
374        """
375        Sets the bio to the SSL pointer which is
376        contained in the connection object.
377        """
378        self._pyfree = 0
379        m2.bio_set_ssl(self.bio, conn.ssl, close_flag)
380        if close_flag == m2.bio_noclose:
381            conn.set_ssl_close_flag(m2.bio_close)
382
383    def do_handshake(self):
384        # type: () -> int
385        """Do the handshake.
386
387        Return 1 if the handshake completes
388        Return 0 or a negative number if there is a problem
389        """
390        return m2.bio_do_handshake(self.bio)
391