Add retain_file_offset helper to utils.py (#21)

dissect/cobaltstrike/utils.py

@@ -10,6 +10,7 @@
 import string
 import itertools
 from functools import partial, wraps
+from contextlib import contextmanager
 
 from typing import BinaryIO, Iterator
 
@@ -24,8 +25,34 @@ def xor(data: bytes, key: bytes) -> bytes:
  return bytes(data)
 
 
+@contextmanager
+def retain_file_offset(fobj, offset=None, whence=io.SEEK_SET):
+ """Return a context manager that changes the position of the file-like object `fobj` to the given byte `offset`.
+ After completion of the block it restores the original position of the file.
+
+ Args:
+ fobj: file-like object
+ offset: offset to seek to relative to position indicated by `whence`. If ``None`` no seek will be done.
+ whence: default is ``SEEK_SET``, values for `whence` are:
+
+ - ``SEEK_SET`` or ``0`` – start of the stream (the default); offset should be zero or positive
+ - ``SEEK_CUR`` or ``1`` – current stream position; offset may be negative
+ - ``SEEK_END`` or ``2`` – end of the stream; offset is usually negative
+
+ Returns:
+ context manager
+ """
+ try:
+ pos = fobj.tell()
+ if offset is not None:
+ fobj.seek(offset, whence)
+ yield fobj
+ finally:
+ fobj.seek(pos)
+
+
 def catch_sigpipe(func):
- """Catches KeyboardInterrupt and BrokenPipeError (OSError 22 on Windows)."""
+ """Decorator for catching KeyboardInterrupt and BrokenPipeError (OSError 22 on Windows)."""
 
  @wraps(func)
  def wrapper(*args, **kwargs):