WebSocket API, token parser updates

tests/js2py_test.py

@@ -3,19 +3,6 @@ import re
 import base64
 import js2py
 
-# Emulate 'atob' function
-# print(base64.standard_b64decode('MmlYaDVXNXVFWXE1ZldKSWF6UTY='))
-
-# Test cases
-# tests = [
-#     """(() => {window[("A" + "J" + "AX_T" + "OKE" + "N")]=("2iXh5W5u" + "EYq" + "5fWJIa" + "zQ6");})();""",
-#     """(() => {window[["N","TOKE","AJAX_"].reverse().join('')]=["IazQ6","fWJ","h5W5uEYq5","2iX"].reverse().join('');})();""",
-#     """(() => {window["AJAX_TOKEN"] = atob("SGVsbG8sIHdvcmxk")})();""",
-#     """(() => {window[atob('QUpBWF9UT0tFTg==')]=atob('MmlYaDVXNXVFWXE1ZldKSWF6UTY=');})();""",
-#     """(() => {window["AJAX_TOKEN"] = "1234" })();""",
-#     """(() => {window[atob('QUpBWF9UT0tFTg==')]="2iXh5W5uEYq5fWJIazQ6";})();""",
-# ]
-
 # Use tests from a file
 tests = []
 with open('../token.txt', 'rt') as f:
@@ -23,10 +10,7 @@ with open('../token.txt', 'rt') as f:
 	del lines[len(lines)-1] # Remove empty string
 	tests = lines
 
-brkregex = re.compile(r'\((?!\)|[\'\"])(.+?)(?<!\(|[\'\"])\)')
-
-def parse_brackets(f):
-	return brkregex.search(f)[1]
+arrowre = re.compile(r'(\w+?|\(\w+?(?:,\s*\w+?)*\)|\(\))\s*=>\s*({\s*[\s\S]+\s*}|[^\r\n]+?(?:;|$))')
 
 def to_ecma5_function(f):
 	# return "(function() { " + f[f.index("{")+1 : f.index("}")] + "})();"
@@ -39,29 +23,28 @@ def atob(s):
 	return base64.standard_b64decode(str(s)).decode('utf-8')
 
 def arrow_conv(f):
-	if '=>' in f:
-		inner = parse_brackets(f)
-		while brkregex.match(inner) != None:
-			inner = parse_brackets(inner)
-
-		func = re.sub(
-			r'(\w+)\s*=>\s*(.+)',
-			r'function(\1){return \2}', inner
-		)
-		start = f.find(inner)
-		end = start + len(inner)
-		f = f[:start] + func + f[end:]
+	m = arrowre.search(f)
+	while m != None:
+		print(f)
+		params = m.group(1).strip('()')
+		body = m.group(2)
+		if body.startswith('{')\
+		and body.endswith('}'):
+			body = body.strip('{}')
+		else:
+			body = f'return {body}'
+		f = arrowre.sub(f'function({params}){{{body}}}', f)
+		m = arrowre.search(f)
+		print(f)
+	#print('function(' + m.group(1).strip("()") + '){return ' + m.group(2) + ';}')
 	return f
 
 ctx = js2py.EvalJs({'atob': atob})
 
 for f in tests:
-	try:
-		c = to_ecma5_function(f)
-		ctx.execute(c)
-		print(ctx.window['AJAX_TOKEN'])
-	except Exception as e:
-		print(c, '\n', e)
+	c = to_ecma5_function(f)
+	ctx.execute(c)
+	print(ctx.window['AJAX_TOKEN'])
 
 # Expected output:
 # 2rKOA1IFdBcHhEM616cb
@@ -76,5 +59,5 @@ for f in tests:
 # 2iXh5W5uEYq5fWJIazQ6
 # CuUcmZ27Fb8bVBNw12Vj
 # YPPe8Ph7vzYaZ9PF9oQP
-# (Note: The last three 
+# (Note: The last four 
 # tokens are different)

tests/js2py_test.py.old

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+import re
+import base64
+import js2py
+
+# Use tests from a file
+tests = []
+with open('../token.txt', 'rt') as f:
+	lines = re.split(r'[\r\n]', f.read())
+	del lines[len(lines)-1] # Remove empty string
+	tests = lines
+
+brkregex = re.compile(r'\((?!\)|[\'\"])(.+?)(?<!\(|[\'\"])\)')
+
+def parse_brackets(f, arrow=True):
+
+	match = brkregex.finditer(f)
+	if not arrow:
+		return match[0].group(1)
+
+	for r in match:
+		func = r.group(1)
+		if '=>' in func:
+			return func
+
+def to_ecma5_function(f):
+	# return "(function() { " + f[f.index("{")+1 : f.index("}")] + "})();"
+	fnstart = f.find('{')+1
+	fnend = f.rfind('}')
+	f = arrow_conv(f[fnstart:fnend])
+	return f
+
+def atob(s):
+	return base64.standard_b64decode(str(s)).decode('utf-8')
+
+def arrow_conv(f):
+	if '=>' in f:
+		inner = parse_brackets(f)
+		print(inner)
+		while brkregex.match(inner) != None:
+			print(inner)
+			inner = parse_brackets(inner)
+
+		func = re.sub(
+			r'(\w+)\s*=>\s*(.+)',
+			r'function(\1){return \2}', inner
+		)
+		start = f.find(inner)
+		end = start + len(inner)
+		f = f[:start] + func + f[end:]
+		print('*r:', f)
+	return f
+
+ctx = js2py.EvalJs({'atob': atob})
+
+for f in tests:
+	c = to_ecma5_function(f)
+	ctx.execute(c)
+	print(ctx.window['AJAX_TOKEN'])
+
+# Expected output:
+# 2rKOA1IFdBcHhEM616cb
+# 2rKOA1IFdBcHhEM616cb
+# 2rKOA1IFdBcHhEM616cb
+# 2rKOA1IFdBcHhEM616cb
+# 2rKOA1IFdBcHhEM616cb
+# 2rKOA1IFdBcHhEM616cb
+# 2rKOA1IFdBcHhEM616cb
+# 2rKOA1IFdBcHhEM616cb
+# 2rKOA1IFdBcHhEM616cb
+# 2iXh5W5uEYq5fWJIazQ6
+# CuUcmZ27Fb8bVBNw12Vj
+# YPPe8Ph7vzYaZ9PF9oQP
+# (Note: The last four 
+# tokens are different)