Skip to content

Commit

Permalink
Allow mode selecting within combined short arguments. (#45)
Browse files Browse the repository at this point in the history
Instead of giving mode separately like `-e -A ...` you can now just use `-eA ...` but only as the first argument.
  • Loading branch information
rocketdey authored Dec 7, 2021
1 parent f397f9c commit 8030c56
Show file tree
Hide file tree
Showing 3 changed files with 246 additions and 10 deletions.
15 changes: 8 additions & 7 deletions covert/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,20 +99,21 @@ def argparse():
sys.exit(0)
ad = {}
args = Args()
# Separate mode selector from other arguments
if av[0].startswith("-") and len(av[0]) > 2:
av.insert(1, f'-{av[0][2:]}')
av[0] = av[0][:2]
# Support a few other forms for Age etc. compatibility (but only as the first arg)
if av[0].lower() in ('enc', 'encrypt', '-e', '--encrypt'):
if av[0] in ('enc', 'encrypt', '-e'):
args.mode, ad, modehelp = 'enc', encargs, f"{hdrhelp}\nEncryption options:\n{enchelp}"
elif av[0].lower() in ('dec', 'decrypt', '-d', '--decrypt'):
elif av[0] in ('dec', 'decrypt', '-d'):
args.mode, ad, modehelp = 'dec', decargs, f"{hdrhelp}\nEncryption options:\n{enchelp}"
elif av[0].lower() in ('bench', 'benchmark'):
elif av[0] in ('bench', 'benchmark'):
args.mode, ad, modehelp = 'benchmark', benchargs, f"{hdrhelp}"
elif next((k for k, v in encargs.items() if av[0] in v), None) is not None:
args.mode, ad, modehelp = 'enc', encargs, enchelp
av.insert(0, None)
else:
stderr.write(' 💣 Invalid or missing command (enc/dec/benchmark).\n')
sys.exit(1)

aiter = iter(av[1:])
longargs = [flag[1:] for switches in ad.values() for flag in switches if flag.startswith("--")]
shortargs = [flag[1:] for switches in ad.values() for flag in switches if not flag.startswith("--")]
Expand Down
218 changes: 218 additions & 0 deletions tests/asdf.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,218 @@
import os
import sys
from sys import stderr, stdout
import colorama
import covert
from covert.cli import main_benchmark, main_dec, main_enc
import mmap
import os
import sys
import itertools
from concurrent.futures import ThreadPoolExecutor
from contextlib import suppress
from io import BytesIO
from pathlib import Path
from time import perf_counter

import pyperclip
from tqdm import tqdm

from covert import lazyexec, passphrase, pubkey, tty, util
from covert.archive import Archive
from covert.blockstream import decrypt_file, encrypt_file

hdrhelp = """\
Usage:
covert enc [files] [recipients] [signatures] [-A | -o unsuspicious.dat [-a]]
covert dec [-A | unsuspicious.dat] [-i id_ed25519] [-o filesfolder]
covert benchmark
Note: covert enc/dec without arguments ask for password and message. Files and
folders get attached together with a message if 'enc -' is specified.
"""

enchelp = """\
-p Passphrase recipient (default)
--wide-open Anyone can open the file (no recipients)
-r PKEY -R FILE Recipient pubkey, .pub file or github:username
-i SKEY Sign with a secret key (string token or id file)
-A Auto copy&paste: ciphertext is copied
-o FILENAME Encrypted file to output (binary unless -a is used)
--pad PERCENT Preferred padding amount (default 5 %)
"""

dechelp = """\
-A Auto copy&paste: ciphertext is pasted
-i SKEY Decrypt with secret key (token or file)
-o FILEFOLDER Extract any attached files to
"""

cmdhelp = f"""\
Covert {covert.__version__} - A file and message encryptor with strong anonymity
💣 Things encrypted with this developer preview mayn't be readable evermore
{hdrhelp}
{enchelp}
{dechelp}
"""


class Args:

def __init__(self):
self.files = []
self.wideopen = None
self.askpass = 0
self.passwords = []
self.recipients = []
self.recipfiles = []
self.outfile = []
self.identities = []
self.padding = "5"
self.armor = None
self.paste = None
self.debug = None


encargs = dict(
askpass='-p --passphrase'.split(),
passwords='--password'.split(),
wideopen='--wide-open'.split(),
recipients='-r --recipient'.split(),
recipfiles='-R --keyfile --recipients-file'.split(),
identities='-i --identity'.split(),
outfile='-o --out --output'.split(),
armor='-a --armor'.split(),
paste='-A'.split(),
padding='--pad --padding'.split(),
debug='--debug'.split(),
)

decargs = dict(
askpass='-p --passphrase'.split(),
passwords='--password'.split(),
identities='-i --identity'.split(),
outfile='-o --out --output'.split(),
paste='-A'.split(),
debug='--debug'.split(),
)

benchargs = dict(debug='--debug'.split(),)


def argparse():
# Custom parsing due to argparse module's limitations
av = sys.argv[1:]
if not av or any(a.lower() in ('-h', '--help') for a in av):
first, rest = cmdhelp.rstrip().split('\n', 1)
if stdout.isatty():
print(f'\x1B[1;44m{first:78}\x1B[0m\n{rest}')
else:
print(f'{first}\n{rest}')
sys.exit(0)
if any(a.lower() in ('-v', '--version') for a in av):
print(cmdhelp.split('\n')[0])
sys.exit(0)
ad = {}
args = Args()
# Separate mode selector from other arguments
if av[0].startswith("-") and len(av[0]) > 2:
av.insert(1, f'-{av[0][2:]}')
av[0] = av[0][:2]
# Support a few other forms for Age etc. compatibility (but only as the first arg)
if av[0] in ('enc', 'encrypt', '-e'):
args.mode, ad, modehelp = 'enc', encargs, f"{hdrhelp}\nEncryption options:\n{enchelp}"
elif av[0] in ('dec', 'decrypt', '-d'):
args.mode, ad, modehelp = 'dec', decargs, f"{hdrhelp}\nEncryption options:\n{enchelp}"
elif av[0] in ('bench', 'benchmark'):
args.mode, ad, modehelp = 'benchmark', benchargs, f"{hdrhelp}"
else:
stderr.write(' 💣 Invalid or missing command (enc/dec/benchmark).\n')
sys.exit(1)

aiter = iter(av[1:])
longargs = [flag[1:] for switches in ad.values() for flag in switches if flag.startswith("--")]
shortargs = [flag[1:] for switches in ad.values() for flag in switches if not flag.startswith("--")]
for a in aiter:
aprint = a
if not a.startswith('-'):
args.files.append(a)
continue
if a == '-':
args.files.append(True)
continue
if a == '--':
args.files += args
break
if a.startswith('--'):
a = a.lower()
if not a.startswith('--') and len(a) > 2:
if any(arg not in shortargs for arg in list(a[1:])):
falseargs = [arg for arg in list(a[1:]) if arg not in shortargs]
stderr.write(f' 💣 {falseargs} is not an argument: covert {args.mode} {a}\n')
sys.exit(1)
a = [f'-{shortarg}' for shortarg in list(a[1:]) if shortarg in shortargs]
if isinstance(a, str):
a = [a]
for i, av in enumerate(a):
argvar = next((k for k, v in ad.items() if av in v), None)
if isinstance(av, int):
continue
if argvar is None:
stderr.write(f'{modehelp}\n 💣 Unknown argument: covert {args.mode} {aprint}\n')
sys.exit(1)
try:
var = getattr(args, argvar)
if isinstance(var, list):
var.append(next(aiter))
elif isinstance(var, str):
setattr(args, argvar, next(aiter))
elif isinstance(var, int):
setattr(args, argvar, var + 1)
else:
setattr(args, argvar, True)
except StopIteration:
stderr.write(f'{modehelp}\n 💣 Argument parameter missing: covert {args.mode} {aprint}\n')
sys.exit(1)

return args

argparse()

def main():
colorama.init()
# CLI argument processing
args = argparse()
print(args.askpass)
if len(args.outfile) > 1:
raise ValueError('Only one output file may be specified')
args.outfile = args.outfile[0] if args.outfile else None

# A quick sanity check, not entirely reliable
if args.outfile in args.files:
raise ValueError('In-place operation is not supported, cannot use the same file as input and output.')

if args.debug:
if args.mode == "enc":
return main_enc(args)
elif args.mode == "dec":
return main_dec(args)
elif args.mode == "benchmark":
return main_benchmark(args)
else:
raise Exception('This should not be reached')
try:
if args.mode == "enc":
return main_enc(args)
elif args.mode == "dec":
return main_dec(args)
elif args.mode == "benchmark":
return main_benchmark(args)
else:
raise Exception('This should not be reached')
except ValueError as e:
stderr.write(f"Error: {e}\n")
except BrokenPipeError:
stderr.write('I/O error (broken pipe)\n')
except KeyboardInterrupt:
stderr.write("Interrupted.\n")
23 changes: 20 additions & 3 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,8 @@
import sys

def test_argparser(capsys):
from covert.__main__ import argparse # Import *after* capsys wraps stdout/stderr
from covert.__main__ import argparse

# Correct but complex arguments
sys.argv = "covert enc --recipient recipient1 -r recipient2 -Arrp recipient3 recipient4".split()
a = argparse()
Expand All @@ -15,10 +15,27 @@ def test_argparser(capsys):
assert not cap.out
assert not cap.err

# Giving mode within combined arguments
sys.argv = "covert -eArrp recipient1 recipient2".split()
a = argparse()
assert a.recipients == ['recipient1', 'recipient2']
assert a.askpass == 1
assert a.paste is True
cap = capsys.readouterr()
assert not cap.out

# Missing argument parameter
sys.argv = "covert enc -Arrp recipient1".split()
with pytest.raises(SystemExit):
argparse()
cap = capsys.readouterr()
assert not cap.out
assert "Argument parameter missing: covert enc -Arrp …" in cap.err
assert "Argument parameter missing: covert enc -Arrp …" in cap.err
# For combined mode
sys.argv = "covert -eArrp recipient1".split()
with pytest.raises(SystemExit):
argparse()
cap = capsys.readouterr()
assert not cap.out
assert "Argument parameter missing: covert enc -Arrp …" in cap.err

0 comments on commit 8030c56

Please sign in to comment.