Extract calls

arpy.py

@@ -60,6 +60,7 @@
 """
 
 import io
+import os
 import struct
 from typing import Optional, List, Dict, BinaryIO, cast, Union
 
@@ -84,6 +85,9 @@ class ArchiveFormatError(Exception):
 class ArchiveAccessError(IOError):
 	""" Raised on problems with accessing the archived files """
 	pass
+class ExtractBreakoutAttempt(IOError):
+	""" Raised on files which would be extracted above specified path """
+	pass
 
 class ArchiveFileHeader(object):
 	""" File header of an archived file, or a special data segment """
@@ -417,6 +421,50 @@ def open(self, name: Union[bytes,ArchiveFileHeader]) -> ArchiveFileData:
 
 		raise ValueError("Can't look up file using type %s, expected bytes or ArchiveFileHeader" % (type(name),))
 
+	def extract(self, member: bytes, path: Optional[Union[str,bytes]]=None) -> None:
+		filename = os.path.basename(member)
+		if path is None:
+			path = os.getcwd()
+
+		if isinstance(path, bytes):
+			filepath = os.path.join(path, filename)
+		else:
+			filepath = os.path.join(path.encode('utf-8'), filename)
+		buf_size = 8*1024
+
+		ar_member = self.open(member)
+		with open(filepath, 'wb') as f:
+			while True:
+				buffer = ar_member.read(buf_size)
+				if not len(buffer):
+					break
+				f.write(buffer)
+
+	def extractall(self, path: Union[str, bytes], members: Optional[List[bytes]]=None) -> None:
+		self.read_all_headers()
+
+		normpath = os.path.normpath(path)
+		if isinstance(normpath, str):
+			normpath = normpath.encode('utf-8')
+
+		if members is None:
+			sources = list(self.archived_files.keys())
+		else:
+			sources = members
+
+		for member in sources:
+			member_dir = os.path.dirname(member)
+			member_name = os.path.basename(member)
+			filepath = os.path.join(normpath, member_dir, member_name)
+			norm_filepath = os.path.normpath(filepath)
+
+			if os.path.commonpath([normpath, norm_filepath]) != normpath:
+				raise ExtractBreakoutAttempt("file %r would be extracted below specified path" % (member,))
+
+			norm_dirpath = os.path.normpath(os.path.join(normpath, member_dir))
+			os.makedirs(norm_dirpath, exist_ok=True)
+			self.extract(member, norm_filepath)
+
 	def __enter__(self):
 		return self
 

test/test_contents.py

@@ -2,6 +2,7 @@
 import unittest
 import os
 import io
+from unittest.mock import patch, mock_open, call
 
 class ArContents(unittest.TestCase):
 	def test_archive_contents(self):
@@ -13,6 +14,44 @@ def test_archive_contents(self):
 		self.assertEqual(b'test_in_file_2\n', f2_contents)
 		ar.close()
 
+	def test_extract(self):
+		m = mock_open()
+		with arpy.Archive(os.path.join(os.path.dirname(__file__), 'contents.ar')) as ar:
+			with patch('arpy.open', m):
+				ar.extract(b'file1', '/foobar')
+
+		m().write.assert_called_once_with(b'test_in_file_1\n')
+		m().__exit__.assert_called_once_with(None, None, None)
+
+	def test_extract_byte_path(self):
+		m = mock_open()
+		with arpy.Archive(os.path.join(os.path.dirname(__file__), 'contents.ar')) as ar:
+			with patch('arpy.open', m):
+				ar.extract(b'file1', b'/foobar')
+
+		m().write.assert_called_once_with(b'test_in_file_1\n')
+		m().__exit__.assert_called_once_with(None, None, None)
+
+	def test_extractall(self):
+		with arpy.Archive(os.path.join(os.path.dirname(__file__), 'contents.ar')) as ar:
+			with patch.object(ar, 'extract') as m_extract:
+				with patch('os.makedirs') as m_makedirs:
+					ar.extractall('/foobar')
+
+		m_extract.assert_has_calls([
+			call(b'file1', b'/foobar/file1'),
+			call(b'file2', b'/foobar/file2'),
+		], any_order=True)
+		m_makedirs.assert_called_with(b'/foobar', exist_ok=True)
+
+	def test_extractall2(self):
+		with arpy.Archive(os.path.join(os.path.dirname(__file__), 'contents.ar')) as ar:
+			with patch.object(ar, 'extract') as m_extract:
+				with patch('os.makedirs') as m_makedirs:
+					ar.extractall('/foobar', [b'file2'])
+
+		m_extract.assert_called_once_with(b'file2', b'/foobar/file2')
+		m_makedirs.assert_called_once_with(b'/foobar', exist_ok=True)
 
 class ArZipLike(unittest.TestCase):
 	def setUp(self):