build-mar/tests/test_writer.py

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