summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--libmat2/torrent.py20
-rw-r--r--tests/test_corrupted_files.py11
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):