security-and-privacy/mat2
security-and-privacy/mat2
1
0
Fork 0
mirror of synced 2024-11-22 09:14:23 +01:00
Code Releases Activity

Simplify how torrent files are handled

Browse Source 
- Rework the testsuite wrt. torrent
- fail at parser's instantiation on corrupted torrent,
  instead of during `get_meta` or `remove_all` call
This commit is contained in:
jvoisin 2018-07-08 13:47:00 +02:00
parent 7ea362d908
commit b5fcddd6a6
2 changed files with 15 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
libmat2/torrent.py
View File

@ -8,13 +8,16 @@ class TorrentParser(abstract.AbstractParser):
mimetypes = {'application/x-bittorrent', }
whitelist = {b'announce', b'announce-list', b'info'}
def __init__(self, filename):
super().__init__(filename)
with open(self.filename, 'rb') as f:
self.dict_repr = _BencodeHandler().bdecode(f.read())
if self.dict_repr is None:
raise ValueError
def get_meta(self) -> Dict[str, str]:
metadata = {}
with open(self.filename, 'rb') as f:
d = _BencodeHandler().bdecode(f.read())
if d is None:
return {'Unknown meta': 'Unable to parse torrent file "%s".' % self.filename}
for k, v in d.items():
for k, v in self.dict_repr.items():
if k not in self.whitelist:
metadata[k.decode('utf-8')] = v
return metadata
@ -22,15 +25,12 @@ class TorrentParser(abstract.AbstractParser):
def remove_all(self) -> bool:
cleaned = dict()
with open(self.filename, 'rb') as f:
d = _BencodeHandler().bdecode(f.read())
if d is None:
return False
for k, v in d.items():
for k, v in self.dict_repr.items():
if k in self.whitelist:
cleaned[k] = v
with open(self.output_filename, 'wb') as f:
f.write(_BencodeHandler().bencode(cleaned))
self.dict_repr = cleaned # since we're stateful
return True

11
tests/test_corrupted_files.py
View File

@ -46,15 +46,14 @@ class TestCorruptedFiles(unittest.TestCase):
def test_torrent(self):
shutil.copy('./tests/data/dirty.png', './tests/data/clean.torrent')
p = torrent.TorrentParser('./tests/data/clean.torrent')
self.assertFalse(p.remove_all())
expected = {'Unknown meta': 'Unable to parse torrent file "./tests/data/clean.torrent".'}
self.assertEqual(p.get_meta(), expected)
with self.assertRaises(ValueError):
torrent.TorrentParser('./tests/data/clean.torrent')
with open("./tests/data/clean.torrent", "a") as f:
f.write("trailing garbage")
p = torrent.TorrentParser('./tests/data/clean.torrent')
self.assertEqual(p.get_meta(), expected)
with self.assertRaises(ValueError):
torrent.TorrentParser('./tests/data/clean.torrent')
os.remove('./tests/data/clean.torrent')
def test_odg(self):