diff options
| -rw-r--r-- | libmat2/torrent.py | 20 | ||||
| -rw-r--r-- | tests/test_corrupted_files.py | 11 |
2 files changed, 15 insertions, 16 deletions
diff --git a/libmat2/torrent.py b/libmat2/torrent.py index ad49f47..d614136 100644 --- a/libmat2/torrent.py +++ b/libmat2/torrent.py | |||
| @@ -8,13 +8,16 @@ class TorrentParser(abstract.AbstractParser): | |||
| 8 | mimetypes = {'application/x-bittorrent', } | 8 | mimetypes = {'application/x-bittorrent', } |
| 9 | whitelist = {b'announce', b'announce-list', b'info'} | 9 | whitelist = {b'announce', b'announce-list', b'info'} |
| 10 | 10 | ||
| 11 | def __init__(self, filename): | ||
| 12 | super().__init__(filename) | ||
| 13 | with open(self.filename, 'rb') as f: | ||
| 14 | self.dict_repr = _BencodeHandler().bdecode(f.read()) | ||
| 15 | if self.dict_repr is None: | ||
| 16 | raise ValueError | ||
| 17 | |||
| 11 | def get_meta(self) -> Dict[str, str]: | 18 | def get_meta(self) -> Dict[str, str]: |
| 12 | metadata = {} | 19 | metadata = {} |
| 13 | with open(self.filename, 'rb') as f: | 20 | for k, v in self.dict_repr.items(): |
| 14 | d = _BencodeHandler().bdecode(f.read()) | ||
| 15 | if d is None: | ||
| 16 | return {'Unknown meta': 'Unable to parse torrent file "%s".' % self.filename} | ||
| 17 | for k, v in d.items(): | ||
| 18 | if k not in self.whitelist: | 21 | if k not in self.whitelist: |
| 19 | metadata[k.decode('utf-8')] = v | 22 | metadata[k.decode('utf-8')] = v |
| 20 | return metadata | 23 | return metadata |
| @@ -22,15 +25,12 @@ class TorrentParser(abstract.AbstractParser): | |||
| 22 | 25 | ||
| 23 | def remove_all(self) -> bool: | 26 | def remove_all(self) -> bool: |
| 24 | cleaned = dict() | 27 | cleaned = dict() |
| 25 | with open(self.filename, 'rb') as f: | 28 | for k, v in self.dict_repr.items(): |
| 26 | d = _BencodeHandler().bdecode(f.read()) | ||
| 27 | if d is None: | ||
| 28 | return False | ||
| 29 | for k, v in d.items(): | ||
| 30 | if k in self.whitelist: | 29 | if k in self.whitelist: |
| 31 | cleaned[k] = v | 30 | cleaned[k] = v |
| 32 | with open(self.output_filename, 'wb') as f: | 31 | with open(self.output_filename, 'wb') as f: |
| 33 | f.write(_BencodeHandler().bencode(cleaned)) | 32 | f.write(_BencodeHandler().bencode(cleaned)) |
| 33 | self.dict_repr = cleaned # since we're stateful | ||
| 34 | return True | 34 | return True |
| 35 | 35 | ||
| 36 | 36 | ||
diff --git a/tests/test_corrupted_files.py b/tests/test_corrupted_files.py index bc5b9d0..776b0e9 100644 --- a/tests/test_corrupted_files.py +++ b/tests/test_corrupted_files.py | |||
| @@ -46,15 +46,14 @@ class TestCorruptedFiles(unittest.TestCase): | |||
| 46 | 46 | ||
| 47 | def test_torrent(self): | 47 | def test_torrent(self): |
| 48 | shutil.copy('./tests/data/dirty.png', './tests/data/clean.torrent') | 48 | shutil.copy('./tests/data/dirty.png', './tests/data/clean.torrent') |
| 49 | p = torrent.TorrentParser('./tests/data/clean.torrent') | 49 | with self.assertRaises(ValueError): |
| 50 | self.assertFalse(p.remove_all()) | 50 | torrent.TorrentParser('./tests/data/clean.torrent') |
| 51 | expected = {'Unknown meta': 'Unable to parse torrent file "./tests/data/clean.torrent".'} | ||
| 52 | self.assertEqual(p.get_meta(), expected) | ||
| 53 | 51 | ||
| 54 | with open("./tests/data/clean.torrent", "a") as f: | 52 | with open("./tests/data/clean.torrent", "a") as f: |
| 55 | f.write("trailing garbage") | 53 | f.write("trailing garbage") |
| 56 | p = torrent.TorrentParser('./tests/data/clean.torrent') | 54 | with self.assertRaises(ValueError): |
| 57 | self.assertEqual(p.get_meta(), expected) | 55 | torrent.TorrentParser('./tests/data/clean.torrent') |
| 56 | |||
| 58 | os.remove('./tests/data/clean.torrent') | 57 | os.remove('./tests/data/clean.torrent') |
| 59 | 58 | ||
| 60 | def test_odg(self): | 59 | def test_odg(self): |
