325 строки
11 KiB
Python
325 строки
11 KiB
Python
# This Source Code Form is subject to the terms of the Mozilla Public
|
|
# License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
import bz2
|
|
|
|
import pytest
|
|
import six
|
|
from mock import patch
|
|
|
|
from mardor.format import extras_header
|
|
|
|
from mardor.reader import MarReader
|
|
from mardor.writer import MarWriter
|
|
from mardor.writer import add_signature_block
|
|
|
|
from mardor.signing import make_hasher
|
|
from mardor.signing import sign_hash
|
|
from mardor.signing import get_publickey
|
|
from mardor.signing import get_privatekey
|
|
|
|
|
|
def test_writer(tmpdir):
|
|
message_p = tmpdir.join('message.txt')
|
|
message_p.write('hello world')
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('wb') as f:
|
|
with MarWriter(f) as m:
|
|
with tmpdir.as_cwd():
|
|
m.add('message.txt')
|
|
|
|
assert mar_p.size() > 0
|
|
|
|
with mar_p.open('rb') as f:
|
|
with MarReader(f) as m:
|
|
assert m.mardata.additional is None
|
|
assert m.mardata.signatures is None
|
|
assert len(m.mardata.index.entries) == 1
|
|
assert m.mardata.index.entries[0].name == 'message.txt'
|
|
m.extract(str(tmpdir.join('extracted')))
|
|
assert (tmpdir.join('extracted', 'message.txt').read('rb') ==
|
|
b'hello world')
|
|
|
|
|
|
def test_writer_adddir(tmpdir):
|
|
tmpdir.mkdir('foo')
|
|
message_p = tmpdir.join('foo', 'message.txt')
|
|
message_p.write('hello world')
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('wb') as f:
|
|
with MarWriter(f) as m:
|
|
with tmpdir.as_cwd():
|
|
m.add('foo')
|
|
|
|
assert mar_p.size() > 0
|
|
|
|
with mar_p.open('rb') as f:
|
|
with MarReader(f) as m:
|
|
assert m.mardata.additional is None
|
|
assert m.mardata.signatures is None
|
|
assert len(m.mardata.index.entries) == 1
|
|
assert m.mardata.index.entries[0].name == 'foo/message.txt'
|
|
m.extract(str(tmpdir.join('extracted')))
|
|
data = tmpdir.join('extracted', 'foo', 'message.txt').read('rb')
|
|
assert data == b'hello world'
|
|
|
|
|
|
def test_writer_uncompressed(tmpdir):
|
|
message_p = tmpdir.join('message.txt')
|
|
message_p.write('hello world')
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('wb') as f:
|
|
with MarWriter(f) as m:
|
|
with tmpdir.as_cwd():
|
|
m.add('message.txt', compress=None)
|
|
|
|
assert mar_p.size() > 0
|
|
|
|
with mar_p.open('rb') as f:
|
|
with MarReader(f) as m:
|
|
assert m.mardata.additional is None
|
|
assert m.mardata.signatures is None
|
|
assert len(m.mardata.index.entries) == 1
|
|
assert m.mardata.index.entries[0].name == 'message.txt'
|
|
m.extract(str(tmpdir.join('extracted')))
|
|
assert (tmpdir.join('extracted', 'message.txt').read('rb') ==
|
|
b'hello world')
|
|
|
|
|
|
def test_writer_compressed(tmpdir):
|
|
message_p = tmpdir.join('message.txt')
|
|
message_p.write('hello world')
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('wb') as f:
|
|
with MarWriter(f) as m:
|
|
with tmpdir.as_cwd():
|
|
m.add('message.txt', compress='bz2')
|
|
|
|
assert mar_p.size() > 0
|
|
|
|
message_compressed = bz2.compress(b'hello world')
|
|
|
|
with mar_p.open('rb') as f:
|
|
with MarReader(f) as m:
|
|
assert m.mardata.additional is None
|
|
assert m.mardata.signatures is None
|
|
assert len(m.mardata.index.entries) == 1
|
|
assert m.mardata.index.entries[0].name == 'message.txt'
|
|
m.extract(str(tmpdir.join('extracted')), decompress=None)
|
|
assert (tmpdir.join('extracted', 'message.txt').read('rb') ==
|
|
message_compressed)
|
|
|
|
|
|
def test_additional(tmpdir):
|
|
message_p = tmpdir.join('message.txt')
|
|
message_p.write('hello world')
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('w+b') as f:
|
|
with MarWriter(f, productversion='99.9', channel='release') as m:
|
|
with tmpdir.as_cwd():
|
|
m.add('message.txt')
|
|
|
|
assert mar_p.size() > 0
|
|
with mar_p.open('rb') as f:
|
|
with MarReader(f) as m:
|
|
assert m.mardata.additional.count == 1
|
|
assert m.mardata.additional.sections[0].productversion == '99.9'
|
|
assert m.mardata.additional.sections[0].channel == 'release'
|
|
assert m.mardata.signatures.count == 0
|
|
assert len(m.mardata.index.entries) == 1
|
|
assert m.mardata.index.entries[0].name == 'message.txt'
|
|
m.extract(str(tmpdir.join('extracted')))
|
|
assert (tmpdir.join('extracted', 'message.txt').read('rb') ==
|
|
b'hello world')
|
|
|
|
|
|
def test_bad_parameters(tmpdir):
|
|
mar_p = tmpdir.join('test.mar')
|
|
f = mar_p.open('w+b')
|
|
with pytest.raises(ValueError):
|
|
MarWriter(f, productversion='foo')
|
|
with pytest.raises(ValueError):
|
|
MarWriter(f, channel='bar')
|
|
with pytest.raises(ValueError):
|
|
MarWriter(f, signing_key='SECRET')
|
|
with pytest.raises(ValueError):
|
|
MarWriter(f, signing_algorithm='crc')
|
|
with pytest.raises(ValueError):
|
|
message_p = tmpdir.join('message.txt')
|
|
message_p.write('hello world')
|
|
with MarWriter(f) as m:
|
|
with tmpdir.as_cwd():
|
|
m.add_file('message.txt', compress='deflate')
|
|
f.close()
|
|
|
|
|
|
@pytest.mark.parametrize('key_size, algo_id', [
|
|
(2048, 'sha1'),
|
|
(4096, 'sha384'),])
|
|
def test_signing(tmpdir, key_size, algo_id, test_keys):
|
|
private_key, public_key = test_keys[key_size]
|
|
|
|
message_p = tmpdir.join('message.txt')
|
|
message_p.write('hello world')
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('w+b') as f:
|
|
with MarWriter(f, signing_key=private_key, channel='release',
|
|
productversion='99.9', signing_algorithm=algo_id) as m:
|
|
with tmpdir.as_cwd():
|
|
m.add('message.txt')
|
|
|
|
assert mar_p.size() > 0
|
|
with mar_p.open('rb') as f:
|
|
with MarReader(f) as m:
|
|
assert m.mardata.additional.count == 1
|
|
assert m.mardata.signatures.count == 1
|
|
assert len(m.mardata.index.entries) == 1
|
|
assert m.mardata.index.entries[0].name == 'message.txt'
|
|
m.extract(str(tmpdir.join('extracted')))
|
|
assert (tmpdir.join('extracted', 'message.txt').read('rb') ==
|
|
b'hello world')
|
|
assert m.verify(public_key)
|
|
|
|
|
|
def test_addfile_as_dir(tmpdir):
|
|
message_p = tmpdir.join('message.txt')
|
|
message_p.write('hello world')
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('wb') as f:
|
|
with MarWriter(f) as m:
|
|
with tmpdir.as_cwd():
|
|
with pytest.raises(ValueError):
|
|
m.add_dir('message.txt', None)
|
|
|
|
|
|
def test_adddir_as_file(tmpdir):
|
|
message_p = tmpdir.join('subdir', 'message.txt')
|
|
tmpdir.join('subdir').mkdir()
|
|
message_p.write('hello world')
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('wb') as f:
|
|
with MarWriter(f) as m:
|
|
with tmpdir.as_cwd():
|
|
with pytest.raises(ValueError):
|
|
m.add_file('subdir', None)
|
|
|
|
|
|
def test_xz_writer(tmpdir):
|
|
message_p = tmpdir.join('message.txt')
|
|
message_p.write('hello world')
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('wb') as f:
|
|
with MarWriter(f) as m:
|
|
with tmpdir.as_cwd():
|
|
m.add('message.txt', compress='xz')
|
|
|
|
assert mar_p.size() > 0
|
|
|
|
with mar_p.open('rb') as f:
|
|
with MarReader(f) as m:
|
|
assert m.mardata.additional is None
|
|
assert m.mardata.signatures is None
|
|
assert len(m.mardata.index.entries) == 1
|
|
assert m.mardata.index.entries[0].name == 'message.txt'
|
|
m.extract(str(tmpdir.join('extracted')))
|
|
assert (tmpdir.join('extracted', 'message.txt').read('rb') ==
|
|
b'hello world')
|
|
|
|
|
|
def test_writer_badmode(tmpdir, test_keys):
|
|
private_key, public_key = test_keys[2048]
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('wb') as f:
|
|
with pytest.raises(ValueError):
|
|
MarWriter(f, signing_key=private_key, channel='release',
|
|
productversion='99.9', signing_algorithm='sha1')
|
|
|
|
|
|
def test_empty_mar(tmpdir):
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('w+b') as f:
|
|
with MarWriter(f) as m:
|
|
pass
|
|
|
|
with mar_p.open('rb') as f:
|
|
with MarReader(f) as m:
|
|
assert len(m.mardata.index.entries) == 0
|
|
assert not m.mardata.signatures
|
|
|
|
|
|
def test_add_signature(tmpdir, mar_cue, test_keys):
|
|
dest_mar = tmpdir.join('test.mar')
|
|
|
|
# Add a dummy signature
|
|
with mar_cue.open('rb') as s, dest_mar.open('w+b') as f:
|
|
add_signature_block(s, f, 'sha384')
|
|
|
|
with mar_cue.open('rb') as s, MarReader(s) as m, dest_mar.open('rb') as f, MarReader(f) as m1:
|
|
assert m.productinfo == m1.productinfo
|
|
assert m.mardata.additional.sections == m1.mardata.additional.sections
|
|
|
|
assert len(m.mardata.index.entries) == len(m1.mardata.index.entries)
|
|
assert m1.mardata.signatures.count == 1
|
|
|
|
hashes = m1.calculate_hashes()
|
|
assert len(hashes) == 1
|
|
assert hashes[0][1][:20] == b"\r\xa9x\x7f#\xf2m\x93a\xcc\xafJ=\x85\xa3Ss\xb43;"
|
|
|
|
|
|
# Now sign the hash using the test keys, and add the signature back into the file
|
|
private_key, public_key = test_keys[4096]
|
|
|
|
sig = sign_hash(private_key, hashes[0][1], 'sha384')
|
|
# Add the signature back into the file
|
|
with mar_cue.open('rb') as s, dest_mar.open('w+b') as f:
|
|
add_signature_block(s, f, 'sha384', sig)
|
|
|
|
with dest_mar.open('rb') as f, MarReader(f) as m1:
|
|
assert m1.verify(public_key)
|
|
|
|
# Assert file contents are the same
|
|
with dest_mar.open('rb') as f, MarReader(f) as m1:
|
|
with MarReader(mar_cue.open('rb')) as m:
|
|
offset_delta = m1.mardata.data_offset - m.mardata.data_offset
|
|
for (e, e1) in zip(m.mardata.index.entries, m1.mardata.index.entries):
|
|
assert e.name == e1.name
|
|
assert e.flags == e1.flags
|
|
assert e.size == e1.size
|
|
assert e.offset == e1.offset - offset_delta
|
|
|
|
s = b''.join(m.extract_entry(e, decompress=None))
|
|
s1 = b''.join(m1.extract_entry(e1, decompress=None))
|
|
assert len(s) == e.size
|
|
assert len(s1) == e1.size
|
|
assert s == s1
|
|
|
|
|
|
def test_padding(tmpdir):
|
|
"""Check that adding a signature preserves the original padding"""
|
|
message_p = tmpdir.join('message.txt')
|
|
message_p.write('hello world')
|
|
def padded_write(self, productversion, channel):
|
|
self.fileobj.seek(self.additional_offset)
|
|
extras = extras_header.build(dict(
|
|
count=1,
|
|
sections=[dict(
|
|
channel=six.u(channel),
|
|
productversion=six.u(productversion),
|
|
size=len(channel) + len(productversion) + 2 + 8 + 10,
|
|
padding=b'\x00' * 10,
|
|
)],
|
|
))
|
|
self.fileobj.write(extras)
|
|
self.last_offset = self.fileobj.tell()
|
|
|
|
with patch.object(MarWriter, 'write_additional', padded_write):
|
|
mar_p = tmpdir.join('test.mar')
|
|
with mar_p.open('w+b') as f:
|
|
with MarWriter(f, productversion='99.0', channel='1') as m:
|
|
with tmpdir.as_cwd():
|
|
m.add('message.txt', compress='bz2')
|
|
|
|
with mar_p.open('rb') as f:
|
|
with MarReader(f) as m:
|
|
assert m.mardata.additional.sections[0].padding == b'\x00' * 10
|