Skip to content

Commit fb292e9

Browse files
dhondtadhondta
committed
Improved guessing and ranking
1 parent 564aa5a commit fb292e9

2 files changed

Lines changed: 72 additions & 70 deletions

File tree

‎codext/__common__.py‎

Lines changed: 65 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -744,6 +744,10 @@ def remove(name):
744744
json.dump(PERS_MACROS, f, indent=2)
745745
except KeyError:
746746
pass
747+
try:
748+
del CODECS_CACHE[name]
749+
except KeyError:
750+
pass
747751
for s in ["En", "De"]:
748752
try:
749753
delattr(builtins, "%s%scodeError" % (name.capitalize(), s))
@@ -864,6 +868,7 @@ def _handle_error(token, position, output="", eename=None):
864868
"""
865869
if errors == "strict":
866870
msg = "'%s' codec can't %scode %s '%s' in %s %d"
871+
token = ensure_str(token)
867872
token = token[:7] + "..." if len(token) > 10 else token
868873
err = getattr(builtins, exc)(msg % (eename or ename, ["en", "de"][decode], kind, token, item, position))
869874
err.output = output
@@ -968,36 +973,37 @@ def __register(search_function):
968973
codecs.register = __register
969974

970975

971-
def search(encoding_regex):
976+
def search(encoding_regex, extended=True):
972977
""" Function similar to lookup but allows to search for an encoding based on a regex instead. It searches this way
973978
into the local registry but also tries a simple lookup with the original lookup function. """
974979
matches = []
975-
for search_function in __codecs_registry:
980+
for search_function in CODECS_OVERWRITTEN + __codecs_registry:
976981
n = search_function.__name__
977982
for name in [n, n.replace("_", "-")]:
978983
if re.search(encoding_regex, name):
979-
matches.append(name)
984+
matches.append(n.replace("_", "-"))
980985
continue
981-
# in some cases, encoding_regex can match a generated string that uses a particular portion of its generating
982-
# pattern ; e.g. we expect encoding_regex="uu_" to find "uu" and "uu_codec" while it can also find "morse" or
983-
# "atbash" very rarely because of their dynamic patterns and the limited number of randomly generated strings
984-
# so, we can use a qualified majority voting to ensure we do not get a "junk" encoding in the list of matches ;
985-
# executing 5 times the string generation for a given codec but adding the codec to the list of matches only
986-
# if we get at least 3 matches ensures that we consider up to 2 failures that could be stochastic, therefore
987-
# drastically decreasing the probability to get a "junk" encoding in the matches list
988-
c = 0
989-
for i in range(5):
990-
for s in generate_strings_from_regex(search_function.__pattern__):
991-
if re.search(encoding_regex, s):
992-
c += 1
986+
if extended:
987+
# in some cases, encoding_regex can match a generated string that uses a particular portion of its
988+
# generating pattern ; e.g. we expect encoding_regex="uu_" to find "uu" and "uu_codec" while it can also
989+
# find "morse" or "atbash" very rarely because of their dynamic patterns and the limited number of randomly
990+
# generated strings
991+
# so, we can use a qualified majority voting to ensure we do not get a "junk" encoding in the list of
992+
# matches ; executing 5 times the string generation for a given codec but adding the codec to the list of
993+
# matches only if we get at least 3 matches ensures that we consider up to 2 failures that could be
994+
# stochastic, therefore drastically decreasing the probability to get a "junk" encoding in the matches list
995+
c = 0
996+
for i in range(5):
997+
for s in generate_strings_from_regex(search_function.__pattern__):
998+
if re.search(encoding_regex, s):
999+
c += 1
1000+
break
1001+
if c >= 3:
1002+
matches.append(n)
9931003
break
994-
if c >= 3:
995-
matches.append(n)
996-
break
9971004
for s, n in ALIASES.items():
9981005
if re.search(encoding_regex, s) or re.search(encoding_regex, n):
9991006
matches.append(n)
1000-
break
10011007
return sorted(list(set(matches)), key=_human_keys)
10021008
codecs.search = search
10031009

@@ -1241,7 +1247,7 @@ def _validate(stop_function, lang_backend="none"):
12411247
stopfunc._validate = _validate
12421248

12431249

1244-
def __guess(prev_input, input, stop_func, depth, max_depth, min_depth, encodings, codecs, result, found=(),
1250+
def __guess(prev_input, input, stop_func, depth, max_depth, min_depth, encodings, result, found=(),
12451251
stop=True, show=False, scoring_heuristic=False, extended=False, debug=False):
12461252
""" Perform a breadth-first tree search using a ranking logic to select and prune the list of codecs. """
12471253
if depth > min_depth and stop_func(input):
@@ -1255,58 +1261,53 @@ def __guess(prev_input, input, stop_func, depth, max_depth, min_depth, encodings
12551261
return
12561262
prev_enc = found[-1] if len(found) > 0 else ""
12571263
e = encodings.get(depth, encodings.get(-1, []))
1258-
for new_input, encoding in __rank(prev_input, input, prev_enc, e, codecs, scoring_heuristic, extended):
1264+
for new_input, encoding in __rank(prev_input, input, prev_enc, e, scoring_heuristic, extended):
12591265
if len(result) > 0 and stop:
12601266
return
12611267
if debug:
12621268
print("[*] Depth %0{}d/%d: %s".format(len(str(max_depth))) % (depth+1, max_depth, encoding))
1263-
__guess(input, new_input, stop_func, depth+1, max_depth, min_depth, encodings, codecs, result,
1264-
found + (encoding, ), stop, show, scoring_heuristic, extended, debug)
1269+
__guess(input, new_input, stop_func, depth+1, max_depth, min_depth, encodings, result, found + (encoding, ),
1270+
stop, show, scoring_heuristic, extended, debug)
12651271

12661272

12671273
def __make_encodings_dict(include, exclude):
12681274
""" Process encodings inclusion and exclusion lists, listing categories and developping codecs' lists of possible
12691275
encoding names. It also creates a cache with the CodecInfo objects for improving performance. """
1270-
codecs = {}
12711276
def _develop(d, keep=True):
12721277
d = d or {}
12731278
for k, v in d.items():
12741279
l, cc = [], [e for e in v if e in CODECS_CATEGORIES]
1275-
for enc in (list_encodings(*cc) if len(cc) > 0 or keep else [] + \
1280+
# list from in-scope categories and then everything that is not a category
1281+
for enc in ((list_encodings(*cc) if len(cc) > 0 or keep else []) + \
12761282
[e for e in v if e not in CODECS_CATEGORIES]):
1277-
try:
1278-
g = lookup(enc, False).parameters['guess']
1279-
except:
1280-
g = [enc]
1281-
if enc in g and not keep: # e.g. "rot-1" => ["rot-1", "rot-2", ...] ; only "rot-1" is to be selected
1283+
g = []
1284+
for e in (search(enc, False) or [enc]):
1285+
try:
1286+
ci = lookup(e, False)
1287+
g.extend(ci.parameters['guess'])
1288+
except:
1289+
pass
1290+
if enc in g: # e.g. "rot-1" => ["rot-1", "rot-2", ...] ; only "rot-1" is to be selected
12821291
l.append(enc)
1283-
else: # e.g. "rot" => ["rot-1", "rot-2", ...] ; all the "rot-N" shall be selected
1292+
else: # e.g. "rot" => ["rot-1", "rot-2", ...] ; all the "rot-N" shall be selected
12841293
l.extend(g)
1285-
d[k] = l
1286-
if keep:
1287-
for e in l:
1288-
# cache newly loaded CodecInfo objects
1289-
ci = lookup(e, False)
1290-
n = ci.name
1291-
if n in CODECS_CACHE:
1292-
ci = CODECS_CACHE[n] # keep the cached object
1293-
else:
1294-
CODECS_CACHE[n] = ci # cache the new object
1295-
codecs[e] = ci
1294+
d[k] = list(set(l))
12961295
return d
12971296
exclude = _develop(exclude, False)
1298-
return {k: [x for x in v if x not in exclude.get(k, [])] for k, v in _develop(include).items()}, codecs
1297+
return {k: [x for x in v if x not in exclude.get(k, [])] for k, v in _develop(include).items()}
12991298

13001299

1301-
def __rank(prev_input, input, prev_encoding, encodings, codecs, heuristic=False, extended=False, yield_score=False):
1300+
def __rank(prev_input, input, prev_encoding, encodings, heuristic=False, extended=False, yield_score=False):
13021301
""" Filter valid encodings and rank them by relevance. """
13031302
ranking = {}
1304-
for encoding in encodings:
1303+
for e in encodings:
13051304
try:
1306-
score, new = __score(prev_input, input, prev_encoding, encoding, codecs.get(encoding), heuristic, extended)
1307-
except TypeError:
1308-
continue
1309-
ranking[encoding] = (score, new)
1305+
codec = CODECS_CACHE[e]
1306+
except KeyError:
1307+
CODECS_CACHE[e] = codec = lookup(e, False)
1308+
t = __score(prev_input, input, prev_encoding, e, codec, heuristic, extended)
1309+
if t:
1310+
ranking[e] = t
13101311
for encoding, result in sorted(ranking.items(), key=lambda x: -x[1][0]):
13111312
yield result if yield_score else result[1], encoding
13121313

@@ -1315,16 +1316,16 @@ class _Text(object):
13151316
__slots__ = ["entropy", "lcharset", "len", "padding", "printables", "text"]
13161317

13171318
def __init__(self, text, pad_char=None):
1318-
self.text = text
1319-
c = text[-1]
1320-
pad_char, last_char = (b(pad_char), c) if isinstance(c, int) else (pad_char, ord(c))
1321-
self.padding = pad_char is not None and last_char == ord(pad_char)
1319+
self.text = ensure_str(text)
1320+
c = self.text[-1]
1321+
pad_char, last_char = (chr(pad_char), chr(c)) if isinstance(c, int) else (pad_char, c)
1322+
self.padding = pad_char is not None and last_char == pad_char
13221323
if self.padding:
13231324
text = text.rstrip(pad_char)
1324-
self.len = len(text)
1325-
self.lcharset = len(set(text))
1326-
self.printables = float(len([c for c in text if (chr(c) if isinstance(c, int) else c) in printable])) / self.len
1327-
self.entropy = entropy(text)
1325+
self.len = len(self.text)
1326+
self.lcharset = len(set(self.text))
1327+
self.printables = float(len([c for c in self.text if c in printable])) / self.len
1328+
self.entropy = entropy(self.text)
13281329

13291330

13301331
def __score(prev_input, input, prev_encoding, encoding, codec, heuristic=False, extended=False):
@@ -1386,13 +1387,14 @@ def __score(prev_input, input, prev_encoding, encoding, codec, heuristic=False,
13861387
except TypeError:
13871388
expf = expf(f)
13881389
if isinstance(expf, (int, float)):
1390+
tmp = expf
13891391
expf = (1/f - .1 <= 1/expf <= 1/f + .1)
13901392
elif isinstance(expf, (tuple, list)) and len(expf) == 2:
13911393
expf = 1/f - expf[1] <= 1/expf[0] <= 1/f + expf[1]
13921394
s += [-1., .1][expf]
13931395
# afterwards, if the input text has an entropy close to the expected one, give a bonus weighted on the
13941396
# number of input characters to take bad entropies of shorter strings into account
1395-
entr = sc.get('entropy', {})
1397+
entr = sc.get('entropy', lambda e: e)
13961398
entr = entr.get(encoding, entr.get('default')) if isinstance(entr, dict) else entr
13971399
if isinstance(entr, type(lambda: None)):
13981400
try: # this case allows to consider the current encoding name from the current codec
@@ -1401,7 +1403,7 @@ def __score(prev_input, input, prev_encoding, encoding, codec, heuristic=False,
14011403
entr = entr(obj.entropy)
14021404
if entr is not None:
14031405
# use a quadratic heuristic to compute a weight for the entropy delta, aligned on (100w,.1) and (200w,1)
1404-
d_entr = min(4e-05 * obj.len**2 - .003 * obj.len, 1) * abs(entr - entropy(new_input))
1406+
d_entr = min(5.958194e-06 * obj.len**2 - .002381 * obj.len, 1) * abs(entr - entropy(new_input))
14051407
if d_entr <= .5:
14061408
s += .5 - d_entr
14071409
# finally, if relevant, apply a custom bonus (e.g. when a regex pattern is matched)
@@ -1475,12 +1477,11 @@ def guess(input, stop_func=stopfunc.default, min_depth=0, max_depth=5, include=N
14751477
if not isinstance(l, dict) or not all(isinstance(k, int) for k in l.keys()):
14761478
raise ValueError("Include argument shall be a list or a dictionary with integer keys")
14771479
# precompute encodings lists per depth and cache the related CodecInfo objects
1478-
encodings, codecs = __make_encodings_dict(include, exclude)
1479-
result = {}
1480+
encodings, result = __make_encodings_dict(include, exclude), {}
14801481
try:
14811482
# breadth-first search
14821483
for d in range(max_depth):
1483-
__guess("", input, stop_func, 0, d+1, min_depth, encodings, codecs, result, tuple(found), stop, show,
1484+
__guess("", input, stop_func, 0, d+1, min_depth, encodings, result, tuple(found), stop, show,
14841485
scoring_heuristic, extended, debug)
14851486
if stop and len(result) > 0:
14861487
break
@@ -1500,9 +1501,8 @@ def rank(input, extended=False, limit=-1, include=None, exclude=None):
15001501
:param include: inclusion list with category, codec or encoding names (nothing means include every encoding)
15011502
:param exclude: exclusion list with category, codec or encoding names (nothing means exclude no encoding)
15021503
"""
1503-
encodings, codecs = __make_encodings_dict({0: include or CODECS_CATEGORIES}, {0: exclude or []})
1504-
r = list(__rank(None, input, "", encodings[0], codecs, True, extended, True))
1505-
CODECS_CACHE = {}
1504+
encodings = __make_encodings_dict({-1: include or CODECS_CATEGORIES}, {-1: exclude or []})
1505+
r = list(__rank(None, input, "", encodings[-1], True, extended, True))
15061506
return r[:limit] if len(r) > 1 else r
15071507
codecs.rank = rank
15081508

‎codext/__init__.py‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -154,12 +154,14 @@ def _format_action_invocation(self, action):
154154
guess.add_argument("-v", "--verbose", action="store_true",
155155
help="show guessing information and steps (default: False)")
156156
rank = sparsers.add_parser("rank", help="rank the most probable encodings based on the given input", **kw)
157-
rank.add_argument("-c", "--codec-categories", nargs="*", action="extend", metavar="CATEGORY",
158-
help="codec categories to be included in the search ; format: string|tuple|list(strings|tuples)")
159-
rank.add_argument("-e", "--exclude-codecs", nargs="*", action="extend", metavar="CODEC",
160-
help="codecs to be explicitely not used ; format: string|tuple|list(strings|tuples)")
157+
rank.add_argument("-e", "--exclude", nargs="*", action="extend", metavar="CAT|COD|ENC",
158+
help="categories, codecs and encodings to be explicitely not used ;\n "
159+
"format: [category|codec|encoding] OR depth:[category|codec|encoding]")
161160
rank.add_argument("-E", "--extended", action="store_true",
162161
help="while using the scoring heuristic, also consider null scores (default: False)")
162+
rank.add_argument("-i", "--include", nargs="*", action="extend", metavar="CAT|COD|ENC",
163+
help="categories, codecs and encodings to be explicitely used ;\n "
164+
"format: [category|codec|encoding] OR depth:[category|codec|encoding]")
163165
rank.add_argument("-l", "--limit", type=int, default=10, help="limit the number of displayed results")
164166
search = sparsers.add_parser("search", help="search for codecs")
165167
search.add_argument("pattern", nargs="+", help="encoding pattern to search")
@@ -243,7 +245,7 @@ def _format_action_invocation(self, action):
243245
if len(r) == 0:
244246
print("Could not decode :-(")
245247
elif args.command == "rank":
246-
for i, e in codecs.rank(c, args.extended, args.limit, args.codec_categories, args.exclude_codecs):
248+
for i, e in codecs.rank(c, args.extended, args.limit, args.include, args.exclude):
247249
s = "[+] %.5f: %s" % (i[0], e)
248250
print(s if len(s) <= 80 else s[:77] + "...")
249251
except Exception as e:

0 commit comments

Comments
 (0)