Skip to content

Commit

Permalink
Fix bugs in Update Metadata, restore original PPATH_HEAD
Browse files Browse the repository at this point in the history
  • Loading branch information
NotStatilko committed May 15, 2023
1 parent f8bcbc9 commit dd845a9
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 15 deletions.
48 changes: 38 additions & 10 deletions tgbox/api/local.py
Original file line number Diff line number Diff line change
Expand Up @@ -2453,6 +2453,7 @@ def __init__(
logger.debug('Decrypting efile_path with the MainKey')
self._file_path = AES(self._mainkey).decrypt(self._efile_path) # pylint: disable=no-member
self._file_path = Path(self._file_path.decode())
self._original_file_path = self._file_path
else:
logger.warning(
f'''We can\'t decrypt real file path of ID{self._id} because '''
Expand Down Expand Up @@ -2600,7 +2601,7 @@ async def refresh_metadata(
metadata. You can use it as alternative to
the "drb" argument here.
_updated_metadata (``bytes``, optional):
_updated_metadata (``str``, optional):
Updated metadata by itself. This is for
internal use, specify only ``drb``.
Expand Down Expand Up @@ -2634,6 +2635,21 @@ async def refresh_metadata(
# all updated metadata from the LocalBox
_updated_metadata, updates = None, {}

# We also restore original PPATH_HEAD
self._file_path = self._original_file_path
#
if isinstance(self._lb, DecryptedLocalBox):
self._directory = await self._lb._make_local_path(self._file_path)

await self._lb._tgbox_db.FILES.execute((
'UPDATE FILES SET PPATH_HEAD=? WHERE ID=?',
(self._directory.part_id, self._id)
))
else:
logger.warning(
f'''We can not restore the original PPATH_HEAD of the ID{self._id} '''
'''because it wasn\'t decrypted with the DecryptedLocalBox.'''
)
# =============================================== #

logger.debug(
Expand All @@ -2655,7 +2671,7 @@ async def refresh_metadata(
mainkey = self._mainkey if self._mainkey else self._lb._mainkey
self._file_path = Path(AES(mainkey).decrypt(v).decode())
else:
logger.debug(
logger.warning(
'''Updated metadata contains efile_path, however, '''
'''DecryptedLocalBoxFile that you trying to update '''
'''doesn\'t have a MainKey and wasn\'t decrypted with '''
Expand Down Expand Up @@ -2739,9 +2755,9 @@ async def update_metadata(
if 'efile_path' in changes:
raise ValueError('The "changes" should not contain efile_path')

changes = changes.copy()
current_changes = changes.copy()

logger.debug(f'Applying changes {changes} to the ID{self._id}...')
logger.debug(f'Applying changes {current_changes} to the ID{self._id}...')

dlb = dlb if dlb else self._lb
try:
Expand All @@ -2754,21 +2770,33 @@ async def update_metadata(
except (ValueError, TypeError):
updates = {}

new_file_path = changes.pop('file_path', None)
new_file_path = current_changes.pop('file_path', None)
if isinstance(new_file_path, bytes):
new_file_path = new_file_path.decode()

if new_file_path:
directory = await dlb._make_local_path(Path(new_file_path))
self._directory = await dlb._make_local_path(Path(new_file_path))

await dlb._tgbox_db.FILES.execute((
'UPDATE FILES SET PPATH_HEAD=? WHERE ID=?',
(directory.part_id, self._id)
(self._directory.part_id, self._id)
))
efile_path = AES(dlb._mainkey).encrypt(new_file_path.encode())
changes['efile_path'] = efile_path
current_changes['efile_path'] = efile_path

elif new_file_path is not None:
# User requested us to remove updated file
# path from the LocalBox, so we need to
# restore the original PPATH_HEAD
self._file_path = self._original_file_path
self._directory = await dlb._make_local_path(self._original_file_path)

await dlb._tgbox_db.FILES.execute((
'UPDATE FILES SET PPATH_HEAD=? WHERE ID=?',
(self._directory.part_id, self._id)
))

updates.update(changes)
updates.update(current_changes)

for k,v in tuple(updates.items()):
if not v:
Expand All @@ -2785,7 +2813,7 @@ async def update_metadata(

if drb:
drbf = await drb.get_file(self._id)
await drbf.update_metadata(changes)
await drbf.update_metadata(changes, dlb=dlb)

def get_sharekey(self, reqkey: Optional[RequestKey] = None) -> ShareKey:
"""
Expand Down
10 changes: 5 additions & 5 deletions tgbox/api/remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -2047,17 +2047,17 @@ async def update_metadata(
if 'file_path' in changes and not dlb:
raise ValueError('You can\'t change file_path without specifying dlb!')

changes = changes.copy()
current_changes = changes.copy()

logger.debug(f'Applying changes {changes} to the ID{self._id}...')
logger.debug(f'Applying changes {current_changes} to the ID{self._id}...')
try:
message_caption = urlsafe_b64decode(self._message.message)
updates = AES(self._filekey).decrypt(message_caption)
updates = PackedAttributes.unpack(updates)
except (ValueError, TypeError):
updates = {}

new_file_path = changes.pop('file_path', None)
new_file_path = current_changes.pop('file_path', None)
if isinstance(new_file_path, bytes):
new_file_path = new_file_path.decode()

Expand All @@ -2069,15 +2069,15 @@ async def update_metadata(
(directory.part_id, self._id)
))
efile_path = AES(dlb._mainkey).encrypt(new_file_path.encode())
changes['efile_path'] = efile_path
current_changes['efile_path'] = efile_path

# If new_file_path is empty string then it's should be
# a request to remove updated file_path attribute
# from the RemoteBox file and restore default
if new_file_path is not None:
updates.pop('efile_path', None)

updates.update(changes)
updates.update(current_changes)

for k,v in tuple(updates.items()):
if not v:
Expand Down

0 comments on commit dd845a9

Please sign in to comment.