Skip to content

Commit b931f20

Browse files
dhondtadhondta
committed
Improved the guess feature
1 parent a260fa3 commit b931f20

5 files changed

Lines changed: 245 additions & 52 deletions

File tree

‎codext/__common__.py‎

Lines changed: 68 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,8 @@
3030

3131
__all__ = ["add", "add_map", "b", "clear", "codecs", "decode", "encode", "ensure_str", "examples", "guess",
3232
"generate_strings_from_regex", "get_alphabet_from_mask", "handle_error", "list_categories", "list_encodings",
33-
"lookup", "maketrans", "re", "register", "remove", "reset", "s2i", "search", "stopfunc", "BytesIO", "MASKS",
34-
"PY3"]
33+
"lookup", "maketrans", "rank", "re", "register", "remove", "reset", "s2i", "search", "stopfunc", "BytesIO",
34+
"MASKS", "PY3"]
3535
CODECS_REGISTRY = None
3636
MASKS = {
3737
'a': printable,
@@ -847,31 +847,56 @@ def generate_strings_from_regex(regex, star_plus_max=STAR_PLUS_MAX, repeat_max=R
847847
""")
848848
stopfunc.printables = lambda s: all(c in printable for c in ensure_str(s))
849849
stopfunc.regex = lambda p: lambda s: re.search(p, ensure_str(s)) is not None
850-
stopfunc.text = lambda s: stopfunc.printables(s) and entropy(s) < 4.5
850+
stopfunc.text = lambda s: stopfunc.printables(s) and entropy(s) < 4.4
851+
852+
def _lang(lang):
853+
def _test(s):
854+
if not stopfunc.text(s):
855+
return False
856+
try:
857+
return detect(ensure_str(s)) == lang
858+
except:
859+
return False
860+
return _test
851861

852862
try:
853863
from langdetect import detect, PROFILES_DIRECTORY
854864
for lang in [p.replace("-", "") for p in os.listdir(PROFILES_DIRECTORY)]:
855-
setattr(stopfunc, "lang_%s" % lang, lambda s, l=lang: stopfunc.printables(s) and detect(s) == l)
865+
setattr(stopfunc, "lang_%s" % lang, _lang(lang))
856866
except ImportError:
857867
pass
858868

859869

860870
__flag = lambda x: re.search(r"[Ff][Ll1][Aa4@][Gg96]", x) is not None
861871
def _flag(x):
862-
try:
863-
return __flag(ensure_str(b(x).decode("utf16")))
864-
except (UnicodeDecodeError, UnicodeEncodeError):
865-
return __flag(x)
872+
return __flag(ensure_str(x))
866873
stopfunc.flag = _flag
867874

868875

869-
def __guess(prev_input, input, stop_func, depth, max_depth, codec_categories, exclude, result, found=(), stop=True,
870-
show=False, scoring_heuristic=False, extended=False, debug=False):
876+
def __develop(encodings):
877+
""" Private method for developing the input list of encodings, trying to extend it with every encoding name. """
878+
enc = []
879+
for e in (encodings or []):
880+
try:
881+
ci = lookup(e)
882+
g = ci.parameters['guess']
883+
except:
884+
g = [e]
885+
if e in g: # e.g. "rot-1" => ["rot-1", "rot-2", ...] ; only "rot-1" is to be selected
886+
enc.append(e)
887+
else: # e.g. "rot" => ["rot-1", "rot-2", ...] ; all the "rot-N" shall be selected
888+
enc.extend(g)
889+
return enc
890+
891+
892+
def __guess(prev_input, input, stop_func, depth, max_depth, min_depth, codec_categories, exclude, result, found=(),
893+
stop=True, show=False, scoring_heuristic=False, extended=False, debug=False):
871894
""" Perform a breadth-first tree search using a ranking logic to select and prune the list of codecs. """
872-
if depth > 0 and stop_func(input):
895+
if depth > min_depth and stop_func(input):
873896
if not stop and show and found not in result:
874-
s = "[+] %s: %s" % (", ".join(found), ensure_str(input))
897+
s = repr(input)
898+
s = s[2:-1] if s.startswith("b'") and s.endswith("'") else s
899+
s = "[+] %s: %s" % (", ".join(found), s)
875900
print(s if len(s) <= 80 else s[:77] + "...")
876901
result[found] = input
877902
if depth >= max_depth or len(result) > 0 and stop:
@@ -898,26 +923,26 @@ def expand(items, descr=None, transform=None):
898923
raise ValueError("Bad %sformat %s" % (["%s " % descr, ""][descr is None], items))
899924
return r if transform is None else transform(*r)
900925
# parse valid encodings, expanding included/excluded codecs
901-
c, e = expand(codec_categories, "codec_categories", list_encodings), expand(exclude, "exclude")
926+
c, e = expand(codec_categories, "codec_categories", list_encodings), __develop(expand(exclude, "exclude"))
902927
for new_input, encoding in __rank(prev_input, input, c, scoring_heuristic, extended):
903928
if len(result) > 0 and stop:
904929
return
905930
if encoding in e:
906931
continue
907932
if debug:
908933
print("[*] Depth %d/%d ; trying %s" % (depth+1, max_depth, encoding))
909-
__guess(input, new_input, stop_func, depth+1, max_depth, codec_categories, exclude, result,
934+
__guess(input, new_input, stop_func, depth+1, max_depth, min_depth, codec_categories, exclude, result,
910935
found + (encoding, ), stop, show, scoring_heuristic, extended, debug)
911936

912937

913-
def __rank(prev_input, input, codecs, heuristic=False, extended=False):
938+
def __rank(prev_input, input, codecs, heuristic=False, extended=False, yield_score=False):
914939
""" Filter valid encodings and rank them by relevance. """
915940
ranking = {}
916941
for codec in codecs:
917942
for score, new_input, encoding in __score(prev_input, input, codec, heuristic, extended):
918943
ranking[encoding] = (score, new_input)
919944
for encoding, result in sorted(ranking.items(), key=lambda x: -x[1][0]):
920-
yield result[1], encoding
945+
yield result if yield_score else result[1], encoding
921946

922947

923948
class _Text(object):
@@ -939,7 +964,7 @@ def __score(prev_input, input, codec, heuristic=False, extended=False):
939964
except:
940965
continue
941966
# ignore encodings that give an output identical to the input (identity transformation) or to the previous input
942-
if b(input) == b(new_input) or b(prev_input) == b(new_input):
967+
if prev_input is not None and b(input) == b(new_input) or b(prev_input) == b(new_input):
943968
continue
944969
# compute input's characteristics only once and only if the control flow reaches this point
945970
pad = ci.parameters.get('scoring', {}).get('padding_char')
@@ -948,14 +973,14 @@ def __score(prev_input, input, codec, heuristic=False, extended=False):
948973
if heuristic:
949974
# from here, the goal (e.g. if the input is Base32) is to rank candidate encodings (e.g. multiple base
950975
# codecs) so that we can put the right one as early as possible and eventually exclude bad candidates
951-
s = .0
976+
s = -ci.parameters.get('penalty', .0)
952977
# first, apply a bonus if the length of input text's charset is exactly the same as encoding's charset ;
953978
# on the contrary, if the length of input text's charset is strictly greater, give a penalty
954979
lcs = ci.parameters.get('scoring', {}).get('len_charset', 256)
955980
if isinstance(lcs, type(lambda: None)):
956981
lcs = int(lcs(encoding))
957-
if (pad and obj.padding and lcs + 1 == obj.lcharset) or lcs == obj.lcharset:
958-
s += .3
982+
if (pad and obj.padding and lcs + 1 >= obj.lcharset) or lcs >= obj.lcharset:
983+
s += max(.0, round(.6 * (.99 ** (lcs - obj.lcharset)), 5) - .1)
959984
elif (pad and obj.padding and lcs + 1 < obj.lcharset) or lcs < obj.lcharset:
960985
s -= .2 # this can occur for encodings with no_error set to True
961986
# then, take padding into account, giving a bonus if padding is to be encountered and effectively present,
@@ -1000,7 +1025,7 @@ def __score(prev_input, input, codec, heuristic=False, extended=False):
10001025
yield s, new_input, encoding
10011026

10021027

1003-
def guess(input, stop_func=stopfunc.printables, max_depth=5, codec_categories=None, exclude=None, result=None, found=(),
1028+
def guess(input, stop_func=stopfunc.printables, min_depth=0, max_depth=5, codec_categories=None, exclude=None, found=(),
10041029
stop=True, show=False, scoring_heuristic=False, extended=False, debug=False):
10051030
""" Try decoding without the knowledge of the encoding(s). """
10061031
if max_depth <= 0:
@@ -1011,13 +1036,28 @@ def guess(input, stop_func=stopfunc.printables, max_depth=5, codec_categories=No
10111036
if isinstance(stop_func, string_types):
10121037
stop_func = stopfunc.regex(stop_func)
10131038
if len(input) > 0:
1014-
result = result or {}
1015-
# breadth-first search
1016-
for d in range(max_depth):
1017-
__guess("", input, stop_func, 0, d+1, codec_categories, exclude, result, tuple(found), stop, show,
1018-
scoring_heuristic, extended, debug)
1019-
if stop and len(result) > 0:
1020-
return result
1039+
result = {}
1040+
try:
1041+
# breadth-first search
1042+
for d in range(max_depth):
1043+
__guess("", input, stop_func, 0, d+1, min_depth, codec_categories, exclude, result, tuple(found), stop,
1044+
show, scoring_heuristic, extended, debug)
1045+
if stop and len(result) > 0:
1046+
return result
1047+
except KeyboardInterrupt:
1048+
pass
10211049
return result
10221050
codecs.guess = guess
10231051

1052+
1053+
def rank(input, extended=False, limit=-1, codec_categories=None, exclude=None):
1054+
""" Rank the most probable encodings based on the given input. """
1055+
codecs = list_encodings(*(codec_categories or ()))
1056+
for e in __develop(exclude):
1057+
try:
1058+
codecs.remove(e)
1059+
except ValueError:
1060+
pass
1061+
return list(__rank(None, input, codecs, True, extended, True))[:limit]
1062+
codecs.rank = rank
1063+

‎codext/__init__.py‎

Lines changed: 35 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -80,12 +80,14 @@ def main():
8080
guess = sparsers.add_parser("guess", help="try guessing the decoding codecs")
8181
guess.add_argument("encoding", nargs="*", help="list of known encodings to apply (default: none)")
8282
guess.add_argument("-c", "--codec-categories", help="codec categories to be included in the search ; "
83-
"format: string|tuple|list(strings|tuples)")
84-
guess.add_argument("-d", "--depth", default=5, type=int, help="maximum codec search depth (default: 3)")
83+
"format: string|tuple|list(strings|tuples)")
8584
guess.add_argument("-e", "--exclude-codecs", help="codecs to be explicitely not used ; "
86-
"format: string|tuple|list(strings|tuples)")
87-
guess.add_argument("-f", "--stop-function", default="test", help="result checking function (default: text) ; "
85+
"format: string|tuple|list(strings|tuples)")
86+
guess.add_argument("-f", "--stop-function", default="text", help="result checking function (default: text) ; "
8887
"format: printables|text|flag|lang_[bigram]|[regex]")
88+
guess.add_argument("--max-depth", default=5, type=int, help="maximum codec search depth (default: 5)")
89+
guess.add_argument("--min-depth", default=0, type=int, help="minimum codec search depth before triggering results "
90+
"(default: 0)")
8991
guess.add_argument("--extended", action="store_true",
9092
help="while using the scoring heuristic, also consider null scores (default: False)")
9193
guess.add_argument("--heuristic", action="store_true",
@@ -94,6 +96,14 @@ def main():
9496
help="do not stop if a valid output is found (default: False)")
9597
guess.add_argument("-v", "--verbose", action="store_true",
9698
help="show guessing information and steps (default: False)")
99+
rank = sparsers.add_parser("rank", help="rank the most probable encodings based on the given input")
100+
rank.add_argument("-c", "--codec-categories", help="codec categories to be included in the search ; "
101+
"format: string|tuple|list(strings|tuples)")
102+
rank.add_argument("-e", "--exclude-codecs", help="codecs to be explicitely not used ; "
103+
"format: string|tuple|list(strings|tuples)")
104+
rank.add_argument("-l", "--limit", type=int, default=10, help="limit the number of displayed results")
105+
rank.add_argument("--extended", action="store_true",
106+
help="while using the scoring heuristic, also consider null scores (default: False)")
97107
search = sparsers.add_parser("search", help="search for codecs")
98108
search.add_argument("pattern", nargs="+", help="encoding pattern to search")
99109
args = parser.parse_args()
@@ -112,8 +122,11 @@ def main():
112122
c = b("")
113123
for line in __stdin_pipe():
114124
c += line
125+
# strip only the very last (CR)LF
126+
c = c.rstrip("\r\n") if isinstance(c, str) else c.rstrip(b"\r\n")
127+
# strip any other (CR)LF
115128
if args.strip:
116-
c = re.sub(r"\r?\n", "", c)
129+
c = re.sub(r"\r?\n", "", c) if isinstance(c, str) else c.replace(b"\r\n", b"").replace(b"\n", b"")
117130
if args.command in ["decode", "encode"]:
118131
# encode or decode
119132
for encoding in args.encoding:
@@ -125,14 +138,18 @@ def main():
125138
else:
126139
print(ensure_str(c or "Could not decode :-("), end="")
127140
elif args.command == "guess":
128-
sfunc = getattr(stopfunc, args.stop_function, args.stop_function)
129-
r = {}
130-
try:
131-
codecs.guess(c, sfunc, args.depth, __literal_eval(args.codec_categories),
132-
__literal_eval(args.exclude_codecs), r, args.encoding, not args.do_not_stop, True,
133-
args.heuristic, args.extended, args.verbose)
134-
except KeyboardInterrupt:
135-
pass
141+
r = codecs.guess(c,
142+
getattr(stopfunc, args.stop_function, args.stop_function),
143+
args.min_depth,
144+
args.max_depth,
145+
__literal_eval(args.codec_categories),
146+
__literal_eval(args.exclude_codecs),
147+
args.encoding,
148+
not args.do_not_stop,
149+
True, # show
150+
args.heuristic,
151+
args.extended,
152+
args.verbose)
136153
for i, o in enumerate(r.items()):
137154
e, out = o
138155
if len(e) > 0:
@@ -144,4 +161,9 @@ def main():
144161
print(ensure_str(out))
145162
if len(r) == 0:
146163
print("Could not decode :-(")
164+
elif args.command == "rank":
165+
for i, e in codecs.rank(c, args.extended, args.limit,
166+
__literal_eval(args.codec_categories), __literal_eval(args.exclude_codecs)):
167+
s = "[+] %.5f: %s" % (i[0], e)
168+
print(s if len(s) <= 80 else s[:77] + "...")
147169

‎codext/base/_base.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,8 @@ def _decode(input, errors="strict"):
140140

141141
kwargs['len_charset'] = n
142142
kwargs['printables_rate'] = 1.
143-
add("base{}".format(n) if name is None else name, encode, decode, pattern, entropy=nb, **kwargs)
143+
n = "base{}".format(n) if name is None else name
144+
add(n, encode, decode, pattern, entropy=nb, guess=[n], **kwargs)
144145

145146

146147
def base_generic():
@@ -159,5 +160,5 @@ def _decode(input, errors="strict"):
159160

160161
add("base", encode, decode, r"^base[-_]?([2-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])(?:[-_]generic)?$",
161162
guess=["base%d-generic" % i for i in range(2, 255)], entropy=lambda e, n: log(int(n.split("-")[0][4:]), 2),
162-
len_charset=lambda n: int(n.split("-")[0][4:]), printables_rate=1., category="base-generic")
163+
len_charset=lambda n: int(n.split("-")[0][4:]), printables_rate=1., category="base-generic", penalty=.4)
163164

0 commit comments

Comments
 (0)