Skip to content

Commit 68311d2

Browse files
dhondtadhondta
committed
Improved guess-decode feature
1 parent 8271a02 commit 68311d2

4 files changed

Lines changed: 106 additions & 87 deletions

File tree

‎.coveragerc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ exclude_lines =
1010
if.*?__name__.*?==.*?.__main__.:
1111
def main\(\)\:
1212
def __stdin_pipe\(\)\:
13+
def __literal_eval\(o\)\:
1314
except ImportError:
1415
except NameError:
1516
raise NotImplementedError

‎codext/__common__.py‎

Lines changed: 50 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -57,9 +57,8 @@
5757

5858

5959
def add(ename, encode=None, decode=None, pattern=None, text=True, add_to_codecs=False, **kwargs):
60-
"""
61-
This adds a new codec to the codecs module setting its encode and/or decode functions, eventually dynamically naming
62-
the encoding with a pattern and with file handling.
60+
""" This adds a new codec to the codecs module setting its encode and/or decode functions, eventually dynamically
61+
naming the encoding with a pattern and with file handling.
6362
6463
:param ename: encoding name
6564
:param encode: encoding function or None
@@ -165,10 +164,9 @@ class StreamReader(Codec, codecs.StreamReader):
165164

166165

167166
def add_map(ename, encmap, repl_char="?", sep="", ignore_case=None, no_error=False, intype=None, outype=None, **kwargs):
168-
"""
169-
This adds a new mapping codec (that is, declarable with a simple character mapping dictionary) to the codecs module
170-
dynamically setting its encode and/or decode functions, eventually dynamically naming the encoding with a pattern
171-
and with file handling (if text is True).
167+
""" This adds a new mapping codec (that is, declarable with a simple character mapping dictionary) to the codecs
168+
module dynamically setting its encode and/or decode functions, eventually dynamically naming the encoding with
169+
a pattern and with file handling (if text is True).
172170
173171
:param ename: encoding name
174172
:param encmap: characters encoding map ; can be a dictionary of encoding maps (for use with the first capture
@@ -197,11 +195,10 @@ def add_map(ename, encmap, repl_char="?", sep="", ignore_case=None, no_error=Fal
197195

198196
def __generic_code(decode=False):
199197
def _wrapper(param):
200-
"""
201-
The parameter for wrapping comes from the encoding regex pattern ; e.g.
202-
[no pattern] => param will be None everytime
203-
r"barbie[-_]?([1-4])$" => param could be int 1, 2, 3 or 4
204-
r"^morse(|[-_]?.{3})$" => param could be None, "-ABC" (for mapping to ".-/")
198+
""" The parameter for wrapping comes from the encoding regex pattern ; e.g.
199+
[no pattern] => param will be None everytime
200+
r"barbie[-_]?([1-4])$" => param could be int 1, 2, 3 or 4
201+
r"^morse(|[-_]?.{3})$" => param could be None, "-ABC" (for mapping to ".-/")
205202
206203
In order of precedence:
207204
1. when param is a key in mapdict or mapdict is a list of encoding maps (hence in the case of "barbie...",
@@ -544,7 +541,7 @@ def ensure_str(s, encoding='utf-8', errors='strict'):
544541
# make conversion functions compatible with input/output strings/bytes
545542
def fix_inout_formats(f):
546543
""" This decorator ensures that the first output of f will have the same text format as the first input (str or
547-
bytes). """
544+
bytes). """
548545
@wraps(f)
549546
def _wrapper(*args, **kwargs):
550547
a0 = args[0]
@@ -562,7 +559,7 @@ def _wrapper(*args, **kwargs):
562559
# alphabet generation function from a given mask
563560
def get_alphabet_from_mask(mask):
564561
""" This function generates an alphabet from the given mask. The style used is similar to Hashcat ; group keys are
565-
marked with a heading "?". """
562+
marked with a heading "?". """
566563
i, alphabet = 0, ""
567564
while i < len(mask):
568565
c = mask[i]
@@ -579,8 +576,7 @@ def get_alphabet_from_mask(mask):
579576

580577
# generic error handling function
581578
def handle_error(ename, errors, sep="", repl_char="?", repl_minlen=1, decode=False, item="position"):
582-
"""
583-
This shortcut function allows to handle error modes given some tuning parameters.
579+
""" This shortcut function allows to handle error modes given some tuning parameters.
584580
585581
:param ename: encoding name
586582
:param errors: error handling mode
@@ -597,8 +593,7 @@ def handle_error(ename, errors, sep="", repl_char="?", repl_minlen=1, decode=Fal
597593
exec("class %s(ValueError): pass" % exc, glob)
598594

599595
def _handle_error(token, position):
600-
"""
601-
This handles an encoding/decoding error according to the selected handling mode.
596+
""" This handles an encoding/decoding error according to the selected handling mode.
602597
603598
:param token: input token to be encoded/decoded
604599
:param position: token position index
@@ -661,9 +656,8 @@ def lookup(encoding):
661656

662657

663658
def register(search_function, add_to_codecs=False):
664-
"""
665-
Register function for registering new codecs in the local registry of this module and, if required, in the native
666-
codecs registry (for use with the built-in 'open' function).
659+
""" Register function for registering new codecs in the local registry of this module and, if required, in the
660+
native codecs registry (for use with the built-in 'open' function).
667661
668662
:param search_function: search function for the codecs registry
669663
:param add_to_codecs: also add the search function to the native registry
@@ -862,36 +856,43 @@ def _flag(x):
862856
stopfunc.flag = _flag
863857

864858

865-
def __guess(input, stop_func, depth, max_depth, codecs, result, found=(), exclude=()):
859+
def __guess(input, stop_func, depth, max_depth, codec_categories, exclude, result, found=(), stop=True, show=False):
866860
""" Perform a breadth-first tree search using a ranking logic to select and prune the list of codecs. """
867861
if depth > 0 and stop_func(input):
862+
if not stop and show:
863+
s = "[*] %s: %s" % (", ".join(found), ensure_str(input))
864+
print(s if len(s) <= 80 else s[:77] + "...")
868865
result.append((input, found))
869866
return
870867
if depth >= max_depth or len(result) > 0:
871868
return
872-
# format 1: when 'exclude' is a string, take it as the only encoding to be excluded at any depth
873-
if isinstance(exclude, string_types):
874-
e = (exclude, )
875-
# format 2: when 'exclude' is a tuple, consider it as the encodings to be excluded at any depth
876-
elif isinstance(exclude, tuple):
877-
e = exclude
878-
# format 3: when 'exclude' is a list, consider it as the list of tuples of encodings to be exlucded with the order
879-
# number corresponding to the applicable depth
880-
elif isinstance(exclude, list):
881-
try:
882-
e = exclude[depth]
883-
if e is None:
884-
e = ()
885-
elif isinstance(e, string_types):
886-
e = (e, )
887-
except IndexError:
888-
e = ()
889-
else:
890-
raise ValueError("Bad exclude format %s" % exclude)
891-
for new_input, encoding in __rank(input, codecs):
869+
# compute included and excluded codecs for this depth
870+
def __expand(items, descr=None, transform=None):
871+
# format 1: when string, take it as the only items at any depth
872+
if isinstance(items, string_types):
873+
r = (items, )
874+
# format 2: when tuple, consider it as a list of items at any depth
875+
elif isinstance(items, tuple):
876+
r = items
877+
# format 3: when list, consider it as the list of tuples of items with the order number corresponding to the
878+
# applicable depth
879+
elif isinstance(items, list):
880+
try:
881+
r = items[depth] or ()
882+
if isinstance(r, string_types):
883+
r = (r, )
884+
except IndexError:
885+
r = ()
886+
else:
887+
raise ValueError("Bad %sformat %s" % (["%s " % descr, ""][descr is None], items))
888+
return r if transform is None else transform(*r)
889+
# parse valid encodings, expanding included/excluded codecs
890+
c, e = __expand(codec_categories, "codec_categories", list_encodings), __expand(exclude, "exclude")
891+
for new_input, encoding in __rank(input, c):
892892
if encoding in e:
893893
continue
894-
__guess(new_input, stop_func, depth+1, max_depth, codecs, result, found + (encoding, ))
894+
__guess(new_input, stop_func, depth+1, max_depth, codec_categories, exclude, result, found + (encoding, ), stop,
895+
show)
895896

896897

897898
def __rank(input, codecs):
@@ -924,7 +925,8 @@ def __score(input, codec):
924925
yield score, new_input, encoding
925926

926927

927-
def guess(input, stop_func=stopfunc.printables, max_depth=5, codec_categories=None, found=(), exclude=None):
928+
def guess(input, stop_func=stopfunc.printables, max_depth=5, codec_categories=None, exclude=None, found=(), stop=True,
929+
show=False):
928930
""" Try decoding without the knowledge of the encoding(s). """
929931
if max_depth <= 0:
930932
raise ValueError("Depth must be a non-null positive integer")
@@ -934,18 +936,12 @@ def guess(input, stop_func=stopfunc.printables, max_depth=5, codec_categories=No
934936
if isinstance(stop_func, string_types):
935937
p = stop_func
936938
stop_func = lambda s: re.search(ensure_str(p).lower(), ensure_str(s).lower()) is not None
937-
if codec_categories is None:
938-
codecs = list_encodings()
939-
elif isinstance(codec_categories, string_types):
940-
codecs = list_encodings(codec_categories)
941-
elif isinstance(codec_categories, (tuple, list, set)):
942-
codecs = list_encodings(*codec_categories)
943939
if len(input) > 0:
944940
result = []
945941
for d in range(max_depth):
946-
__guess(input, stop_func, 0, d+1, codecs, result, tuple(found), exclude or [])
947-
if len(result) > 0:
948-
return result[0]
949-
return None, None
942+
__guess(input, stop_func, 0, d+1, codec_categories or [], exclude or [], result, tuple(found), stop, show)
943+
if stop and len(result) > 0:
944+
return result
945+
return result
950946
codecs.guess = guess
951947

‎codext/__init__.py‎

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
44
"""
55
from __future__ import print_function
6+
from ast import literal_eval
67
from six import b, binary_type, text_type
78

89
from .__common__ import *
@@ -23,6 +24,14 @@
2324
reset()
2425

2526

27+
def __literal_eval(o):
28+
""" Non-failing ast.literal_eval alias function. """
29+
try:
30+
return literal_eval(str(o))
31+
except ValueError:
32+
return literal_eval("'" + str(o) + "'")
33+
34+
2635
def __stdin_pipe():
2736
""" Stdin pipe read function. """
2837
try:
@@ -69,8 +78,12 @@ def main():
6978
help="error handling")
7079
guess = sparsers.add_parser("guess", help="try guessing the decoding codecs")
7180
guess.add_argument("encoding", nargs="*", help="list of known encodings to apply")
72-
guess.add_argument("-c", "--category", choices=list_categories(), nargs="*", help="codec categories to search in")
81+
guess.add_argument("-c", "--codec-categories", help="codec categories to be included in the search\n"
82+
"format: string|tuple|list(strings|tuples) ; see the documentation")
7383
guess.add_argument("-d", "--depth", default=3, type=int, help="maximum codec search depth")
84+
guess.add_argument("-e", "--exclude-codecs", help="codecs to be explicitely not used\n"
85+
"format: string|tuple|list(strings|tuples) ; see the documentation")
86+
guess.add_argument("-s", "--do-not-stop", action="store_true", help="do not stop if a valid output is found")
7487
search = sparsers.add_parser("search", help="search for codecs")
7588
search.add_argument("pattern", nargs="+", help="encoding pattern to search")
7689
args = parser.parse_args()
@@ -95,14 +108,24 @@ def main():
95108
# encode or decode
96109
for encoding in args.encoding:
97110
c = getattr(codecs, ["encode", "decode"][args.command == "decode"])(c, encoding, args.errors)
111+
# handle output file or stdout
112+
if args.outfile:
113+
with open(args.outfile, 'wb') as f:
114+
f.write(c)
115+
else:
116+
print(ensure_str(c or "Could not decode :-("), end="")
98117
elif args.command == "guess":
99-
c, e = codecs.guess(c, max_depth=args.depth, codec_categories=args.category, found=args.encoding)
100-
if len(e) > 0:
101-
print("Codecs: %s" % ", ".join(e))
102-
# handle output file or stdout
103-
if args.outfile:
104-
with open(args.outfile, 'wb') as f:
105-
f.write(c)
106-
else:
107-
print(ensure_str(c or "Could not decode :-("), end="")
118+
l = [o for o in codecs.guess(c, stopfunc.printables, args.depth, __literal_eval(args.codec_categories),
119+
__literal_eval(args.exclude_codecs), args.encoding, not args.do_not_stop, True)]
120+
for i, o in enumerate(l):
121+
out, e = o
122+
if len(e) > 0:
123+
if args.outfile:
124+
n, ext = os.path.splitext(args.outfile)
125+
fn = args.outfile if len(l) == 1 else "%s-%d%s" % (n, i+1, ext)
126+
else:
127+
print("Codecs: %s" % ", ".join(e))
128+
print(ensure_str(out), end="")
129+
if len(l) == 0:
130+
print("Could not decode :-(")
108131

‎tests/test_common.py‎

Lines changed: 22 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -130,20 +130,20 @@ def test_search_codecs(self):
130130
def test_guess_decode(self):
131131
codext.reset()
132132
STR = "This is a test"
133-
self.assertEqual(STR, codext.guess("VGhpcyBpcyBhIHRlc3Q=", "a test", 1)[0])
134-
self.assertEqual(STR, codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", "a test", found=["base62"])[0])
133+
self.assertEqual(STR, codext.guess("VGhpcyBpcyBhIHRlc3Q=", "a test", 1)[0][0])
134+
self.assertEqual(STR, codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", "a test", found=["base62"])[0][0])
135135
if hasattr(codext.stopfunc, "lang_en"):
136136
f = codext.stopfunc.lang_en
137-
self.assertEqual(STR, codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", f, found=["base62"])[0])
138-
self.assertIsNotNone(codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", f, max_depth=1)[0])
139-
self.assertEqual(STR, codext.guess("VGhpcyBpcyBhIHRlc3Q=", "a test", 1, "base", exclude=["base100"])[0])
140-
self.assertEqual(STR, codext.guess("VGhpcyBpcyBhIHRlc3Q=", "a test", 1, ["base", "crypto"])[0])
141-
self.assertIsNone(codext.guess("NOT THE ENCODED TEST STRING", "a test", 1, exclude=[None])[0])
142-
self.assertIn("F1@9", codext.guess("VGVzdCBGMUA5ICE=", codext.stopfunc.flag, 1)[0])
143-
self.assertIsNone(codext.guess("VGhpcyBpcyBhIHRlc3Q=", " a test", 1, codec_categories="base",
144-
exclude="base64")[0])
145-
self.assertIsNone(codext.guess("VGhpcyBpcyBhIHRlc3Q=", " a test", 1, codec_categories="base",
146-
exclude=("base64", "atbash"))[0])
137+
self.assertEqual(STR, codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", f, found=["base62"])[0][0])
138+
self.assertIsNotNone(codext.guess("CJG3Ix8bVcSRMLOqwDUg28aDsT7", f, max_depth=1)[0][0])
139+
self.assertEqual(STR, codext.guess("VGhpcyBpcyBhIHRlc3Q=", "a test", 1, "base", exclude=["base100"])[0][0])
140+
self.assertEqual(STR, codext.guess("VGhpcyBpcyBhIHRlc3Q=", "a test", 1, ["base", "crypto"])[0][0])
141+
self.assertEqual(len(codext.guess("NOT THE ENCODED TEST STRING", "a test", 1, exclude=[None])), 0)
142+
self.assertIn("F1@9", codext.guess("VGVzdCBGMUA5ICE=", codext.stopfunc.flag, 1, stop=False, show=True)[0][0])
143+
self.assertEqual(len(codext.guess("VGhpcyBpcyBhIHRlc3Q=", " a test", 1, codec_categories="base",
144+
exclude="base64")), 0)
145+
self.assertEqual(len(codext.guess("VGhpcyBpcyBhIHRlc3Q=", " a test", 1, codec_categories="base",
146+
exclude=("base64", "atbash"))), 0)
147147
self.assertRaises(ValueError, codext.guess, STR, max_depth=0)
148148
self.assertRaises(ValueError, codext.guess, STR, exclude=42)
149149
for c in ["base", "language", "native", "stegano"]:
@@ -159,15 +159,14 @@ def test_guess_decode(self):
159159
enc = codext.encode(b(STR), encoding)
160160
if codext.decode(enc, encoding) == STR:
161161
continue
162-
found_dec, found_encodings = codext.guess(enc, "a test", 1, [c])
163-
print(encoding, found_encodings)
164-
self.assertEqual(ensure_str(STR).lower(), ensure_str(found_dec).lower())
165-
if c != "base":
166-
# do not check for base as the guessed encoding name can be different, e.g.:
167-
# actual: base2
168-
# guessed: base2-generic
169-
if "-icase" in encoding:
170-
self.assertEqual(encoding.lower(), found_encodings[0].lower())
171-
else:
172-
self.assertEqual(encoding, found_encodings[0])
162+
for found_dec, found_encodings in codext.guess(enc, "a test", 1, [c]):
163+
self.assertEqual(ensure_str(STR).lower(), ensure_str(found_dec).lower())
164+
if c != "base":
165+
# do not check for base as the guessed encoding name can be different, e.g.:
166+
# actual: base2
167+
# guessed: base2-generic
168+
if "-icase" in encoding:
169+
self.assertEqual(encoding.lower(), found_encodings[0].lower())
170+
else:
171+
self.assertEqual(encoding, found_encodings[0])
173172

0 commit comments

Comments
 (0)