def test_linked_crypted_io():
key = os.urandom(32)
buf = os.urandom(1024 * 50 + 14)
raw_len = len(buf)
raw_io = io.BytesIO(buf)
eio = SimpleEncryptIO(raw_io, key, raw_len)
dio = to_decryptio(eio, key)
eio = ChaCha20EncryptIO(dio, key, raw_len)
dio = to_decryptio(eio, key)
eio = AES256CBCEncryptIO(dio, key, raw_len)
dio = to_decryptio(eio, key)
length = 0
dbuf = b""
while True:
d = dio.read(1)
if not d:
break
dbuf += d
assert len(d) == 1
length += 1
assert length == raw_len
assert dbuf == buf
def test_noencryptio():
key = b"123"
buf = os.urandom(1024 * 1024 * 50)
c = io.BytesIO(buf)
enc = c.read()
d = to_decryptio(io.BytesIO(enc), key)
dec = d.read()
assert buf == dec
def test_chacha20encryptio():
key = os.urandom(32)
buf = os.urandom(1024 * 1024 * 50)
bio = io.BytesIO(buf)
c = ChaCha20EncryptIO(bio, key, len(buf))
assert total_len(c) == len(buf) + PADDED_ENCRYPT_HEAD_WITH_SALT_LEN
enc = c.read()
d = to_decryptio(io.BytesIO(enc), key)
assert total_len(d) == len(buf)
dec = d.read()
assert buf == dec
def test_simpleencryptio():
key = "123"
nonce_or_iv = os.urandom(16)
buf = os.urandom(1024 * 1024 * 50)
bio = io.BytesIO(buf)
c = SimpleEncryptIO(bio, key, nonce_or_iv, len(buf))
assert total_len(c) == len(buf) + ENCRYPT_HEAD_LEN
enc = c.read()
d = to_decryptio(io.BytesIO(enc), key)
assert total_len(d) == len(buf)
dec = d.read()
assert buf == dec
def test_aes256cbcencryptio():
key = os.urandom(32)
iv = os.urandom(16)
buf = os.urandom(1024 * 1024 * 50)
bio = io.BytesIO(buf)
c = AES256CBCEncryptIO(bio, key, iv, len(buf))
assert total_len(c) == padding_size(len(buf), 16) + ENCRYPT_HEAD_LEN
enc = c.read()
print("enc", len(enc))
d = to_decryptio(io.BytesIO(enc), key)
# assert total_len(d) == len(buf) # can be wrong
dec = d.read()
print("dec", len(dec))
assert buf == dec
def decrypt_file(from_encrypted: PathLike,
to_decrypted: PathLike,
encrypt_password: bytes = b""):
assert exists(from_encrypted)
dio = to_decryptio(open(from_encrypted, "rb"),
encrypt_password=encrypt_password)
dpath = Path(to_decrypted)
dir_ = dpath.parent
if not dir_.exists():
dir_.mkdir(parents=True)
with dpath.open("wb") as dfd:
while True:
data = dio.read(READ_SIZE)
if not data:
break
dfd.write(data)
def test_aes256cbcencryptio():
key = os.urandom(32)
buf = os.urandom(1024 * 1024 * 50 + 14)
bio = io.BytesIO(buf)
c = AES256CBCEncryptIO(bio, key, len(buf))
assert (total_len(c) == padding_size(len(buf), 16) +
PADDED_ENCRYPT_HEAD_WITH_SALT_LEN)
enc = c.read()
print("enc", len(enc))
dio = to_decryptio(io.BytesIO(enc), key)
# assert total_len(d) == len(buf) # can be wrong
dec = dio.read()
print("dec", len(dec))
assert buf == dec
# Encrypt
# Assert length of Read(size), size > 0
buf = os.urandom(1024 * 50)
bio = io.BytesIO(buf)
c = AES256CBCEncryptIO(bio, key, len(buf))
length = 0
while True:
d = c.read(1)
if not d:
break
assert len(d) == 1
length += 1
assert (total_len(c) == padding_size(len(buf), 16) +
PADDED_ENCRYPT_HEAD_WITH_SALT_LEN)
buf = os.urandom(1024 * 50 + 14)
bio = io.BytesIO(buf)
c = AES256CBCEncryptIO(bio, key, len(buf))
length = 0
while True:
d = c.read(1)
if not d:
break
assert len(d) == 1
length += 1
assert (total_len(c) == padding_size(len(buf), 16) +
PADDED_ENCRYPT_HEAD_WITH_SALT_LEN)
# Decrypt
# Assert length of Read(size), size > 0
buf = os.urandom(1024 * 50)
bio = io.BytesIO(buf)
c = AES256CBCEncryptIO(bio, key, len(buf))
enc = b""
while True:
d = c.read(1)
if not d:
break
enc += d
dio = to_decryptio(io.BytesIO(enc), key)
length = 0
while True:
d = dio.read(1)
if not d:
break
assert len(d) == 1
length += 1
assert length == len(buf)
buf = os.urandom(1024 * 50 + 14)
bio = io.BytesIO(buf)
c = AES256CBCEncryptIO(bio, key, len(buf))
enc = b""
while True:
d = c.read(1)
if not d:
break
enc += d
dio = to_decryptio(io.BytesIO(enc), key)
length = 0
while True:
d = dio.read(1)
if not d:
break
assert len(d) == 1
length += 1
assert length == len(buf)
def download(
self,
url: str,
localpath: str,
cookies: Dict[str, Optional[str]],
downloadparams: DownloadParams = DEFAULT_DOWNLOADPARAMS,
out_cmd: bool = False,
encrypt_password: bytes = b"",
):
global DEFAULT_DOWNLOADER
if not self.which():
self = DEFAULT_DOWNLOADER
localpath_tmp = localpath + ".tmp"
def done_callback(fut: Future):
err = fut.exception()
if not err:
shutil.move(localpath_tmp, localpath)
else:
logger.info("`download`: MeDownloader fails: error: %s", err)
if self == Downloader.me:
self._me_download(
url,
localpath_tmp,
cookies=cookies,
downloadparams=downloadparams,
done_callback=done_callback,
encrypt_password=encrypt_password,
)
return
elif self == Downloader.aget_py:
cmd = self._aget_py_cmd(url, localpath_tmp, cookies,
downloadparams)
elif self == Downloader.aget_rs:
cmd = self._aget_rs_cmd(url, localpath_tmp, cookies,
downloadparams)
elif self == Downloader.aria2:
cmd = self._aria2_cmd(url, localpath_tmp, cookies, downloadparams)
else:
cmd = self._aget_py_cmd(url, localpath_tmp, cookies,
downloadparams)
# Print out command
if out_cmd:
_print(" ".join((repr(c) for c in cmd)))
return
returncode = self.spawn(cmd, downloadparams.quiet)
logger.debug("`download`: cmd returncode: %s", returncode)
if returncode != 0:
print(
f"[italic]{self.value}[/italic] fails. return code: [red]{returncode}[/red]"
)
else:
if encrypt_password:
dio = to_decryptio(open(localpath_tmp, "rb"), encrypt_password)
if isinstance(dio, DecryptIO):
with open(localpath, "wb") as fd:
while True:
buf = dio.read(READ_SIZE)
if not buf:
break
fd.write(buf)
os.remove(localpath_tmp)
return
shutil.move(localpath_tmp, localpath)
