forked from LBRYCommunity/lbry-sdk
136 lines
6.3 KiB
Python
136 lines
6.3 KiB
Python
import os
|
|
import binascii
|
|
import logging
|
|
import typing
|
|
import asyncio
|
|
from lbrynet.blob import MAX_BLOB_SIZE
|
|
from lbrynet.stream.descriptor import StreamDescriptor
|
|
if typing.TYPE_CHECKING:
|
|
from lbrynet.blob.blob_manager import BlobFileManager
|
|
from lbrynet.blob.blob_info import BlobInfo
|
|
from lbrynet.blob.blob_file import BlobFile
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def _get_next_available_file_name(download_directory: str, file_name: str) -> str:
|
|
base_name, ext = os.path.splitext(os.path.basename(file_name))
|
|
i = 0
|
|
while os.path.isfile(os.path.join(download_directory, file_name)):
|
|
i += 1
|
|
file_name = "%s_%i%s" % (base_name, i, ext)
|
|
|
|
return file_name
|
|
|
|
|
|
async def get_next_available_file_name(loop: asyncio.BaseEventLoop, download_directory: str, file_name: str) -> str:
|
|
return await loop.run_in_executor(None, _get_next_available_file_name, download_directory, file_name)
|
|
|
|
|
|
class StreamAssembler:
|
|
def __init__(self, loop: asyncio.BaseEventLoop, blob_manager: 'BlobFileManager', sd_hash: str,
|
|
output_file_name: typing.Optional[str] = None):
|
|
self.output_file_name = output_file_name
|
|
self.loop = loop
|
|
self.blob_manager = blob_manager
|
|
self.sd_hash = sd_hash
|
|
self.sd_blob: 'BlobFile' = None
|
|
self.descriptor: StreamDescriptor = None
|
|
self.got_descriptor = asyncio.Event(loop=self.loop)
|
|
self.wrote_bytes_event = asyncio.Event(loop=self.loop)
|
|
self.stream_finished_event = asyncio.Event(loop=self.loop)
|
|
self.output_path = ''
|
|
self.stream_handle = None
|
|
self.written_bytes: int = 0
|
|
|
|
async def _decrypt_blob(self, blob: 'BlobFile', blob_info: 'BlobInfo', key: str):
|
|
if not blob or not self.stream_handle or self.stream_handle.closed:
|
|
return False
|
|
|
|
def _decrypt_and_write():
|
|
offset = blob_info.blob_num * (MAX_BLOB_SIZE - 1)
|
|
self.stream_handle.seek(offset)
|
|
_decrypted = blob.decrypt(
|
|
binascii.unhexlify(key), binascii.unhexlify(blob_info.iv.encode())
|
|
)
|
|
self.stream_handle.write(_decrypted)
|
|
self.stream_handle.flush()
|
|
self.written_bytes += len(_decrypted)
|
|
log.debug("decrypted %s", blob.blob_hash[:8])
|
|
|
|
await self.loop.run_in_executor(None, _decrypt_and_write)
|
|
return True
|
|
|
|
async def setup(self):
|
|
pass
|
|
|
|
async def after_got_descriptor(self):
|
|
pass
|
|
|
|
async def after_finished(self):
|
|
pass
|
|
|
|
async def assemble_decrypted_stream(self, output_dir: str, output_file_name: typing.Optional[str] = None):
|
|
if not os.path.isdir(output_dir):
|
|
raise OSError(f"output directory does not exist: '{output_dir}' '{output_file_name}'")
|
|
await self.setup()
|
|
self.sd_blob = await self.get_blob(self.sd_hash)
|
|
self.descriptor = await StreamDescriptor.from_stream_descriptor_blob(self.loop, self.blob_manager.blob_dir,
|
|
self.sd_blob)
|
|
await self.after_got_descriptor()
|
|
self.output_file_name = output_file_name or self.descriptor.suggested_file_name
|
|
self.output_file_name = await get_next_available_file_name(self.loop, output_dir, self.output_file_name)
|
|
self.output_path = os.path.join(output_dir, self.output_file_name)
|
|
if not self.got_descriptor.is_set():
|
|
self.got_descriptor.set()
|
|
await self.blob_manager.storage.store_stream(
|
|
self.sd_blob, self.descriptor
|
|
)
|
|
await self.blob_manager.blob_completed(self.sd_blob)
|
|
written_blobs = None
|
|
save_tasks = []
|
|
try:
|
|
with open(self.output_path, 'wb') as stream_handle:
|
|
self.stream_handle = stream_handle
|
|
for i, blob_info in enumerate(self.descriptor.blobs[:-1]):
|
|
if blob_info.blob_num != i:
|
|
log.error("sd blob %s is invalid, cannot assemble stream", self.descriptor.sd_hash)
|
|
return
|
|
while self.stream_handle and not self.stream_handle.closed:
|
|
try:
|
|
blob = await self.get_blob(blob_info.blob_hash, blob_info.length)
|
|
if blob and blob.length != blob_info.length:
|
|
log.warning("Found incomplete, deleting: %s", blob_info.blob_hash)
|
|
await self.blob_manager.delete_blobs([blob_info.blob_hash])
|
|
continue
|
|
if await self._decrypt_blob(blob, blob_info, self.descriptor.key):
|
|
save_tasks.append(asyncio.ensure_future(self.blob_manager.blob_completed(blob)))
|
|
written_blobs = i
|
|
if not self.wrote_bytes_event.is_set():
|
|
self.wrote_bytes_event.set()
|
|
log.debug("written %i/%i", written_blobs, len(self.descriptor.blobs) - 2)
|
|
break
|
|
except FileNotFoundError:
|
|
log.debug("stream assembler stopped")
|
|
return
|
|
except (ValueError, IOError, OSError):
|
|
log.warning("failed to decrypt blob %s for stream %s", blob_info.blob_hash,
|
|
self.descriptor.sd_hash)
|
|
continue
|
|
finally:
|
|
if written_blobs == len(self.descriptor.blobs) - 2:
|
|
log.debug("finished decrypting and assembling stream")
|
|
if save_tasks:
|
|
await asyncio.wait(save_tasks)
|
|
await self.after_finished()
|
|
self.stream_finished_event.set()
|
|
else:
|
|
log.debug("stream decryption and assembly did not finish (%i/%i blobs are done)", written_blobs or 0,
|
|
len(self.descriptor.blobs) - 2)
|
|
if self.output_path and os.path.isfile(self.output_path):
|
|
log.debug("erasing incomplete file assembly: %s", self.output_path)
|
|
os.unlink(self.output_path)
|
|
|
|
async def get_blob(self, blob_hash: str, length: typing.Optional[int] = None) -> 'BlobFile':
|
|
return self.blob_manager.get_blob(blob_hash, length)
|