Skip to content

Commit 1934148

Browse files
GnarLitoGnarLito
committed
http checks/convertors added
1 parent 9369f65 commit 1934148

File tree

8 files changed

+334
-183
lines changed

8 files changed

+334
-183
lines changed

‎tryhackme/checks.py‎

Lines changed: 35 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -1,41 +1,44 @@
1-
21
from .errors import *
32

4-
class _base_check:
5-
@property
6-
def name(self):
7-
return type(self).__name__.split('_')[1]
3+
def check(predicate):
4+
def decorator(func):
5+
if not hasattr(func, '__checks__'):
6+
func.__checks__ = []
7+
func.__checks__.append(predicate)
8+
return func
9+
10+
decorator.predicate = predicate
11+
return decorator
12+
13+
def arg_injector(predicate):
14+
def decorator(func):
15+
if not hasattr(func, '__arg_injectors__'):
16+
func.__arg_injectors__ = []
17+
func.__arg_injectors__.append(predicate)
18+
return func
19+
20+
decorator.predicate = predicate
21+
return decorator
822

923

10-
class _type_check(_base_check):
11-
_TYPES = ()
12-
_DEFAULT = None
13-
def convert(self, cls, arg):
14-
if arg is None: return self._DEFAULT
15-
if arg in self._TYPES: return arg
16-
else: raise TypeNotInTypeList(f"'{str(arg)}' is an invalid {self.name} type.")
24+
# * HTTP checks
1725

18-
class _notNone_check(_base_check):
19-
def convert(self, cls, arg):
20-
if arg is not None:
21-
return arg
22-
else: raise NotValidUrlParameters(f"'{str(arg)}' is not allowed to be null {self.name} type.")
26+
def is_authenticated():
27+
def predicate(cls):
28+
if cls.authenticated:
29+
return True
30+
raise SessionRequired()
31+
return check(predicate)
2332

24-
class _vpn_types(_type_check):
25-
_TYPES = ('machines', 'networks')
26-
_DEFAULT = 'machines'
2733

28-
class _county_types(_type_check):
29-
_TYPES = ("AF","AX","AL","DZ","AS","AD","AO","AI","AQ","AG","AR","AM","AW","AU","AT","AZ","BS","BH","BD","BB","BY","BE","BZ","BJ","BM","BT","BO","BQ","BA","BW","BV","BR","IO","BN","BG","BF","BI","KH","CM","CA","CV","KY","CF","TD","CL","CN","CX","CC","CO","KM","CG","CD","CK","CR","CI","HR","CU","CW","CY","CZ","DK","DJ","DM","DO","EC","EG","SV","GQ","ER","EE","ET","FK","FO","FJ","FI","FR","GF","PF","TF","GA","GM","GE","DE","GH","GI","GR","GL","GD","GP","GU","GT","GG","GN","GW","GY","HT","HM","VA","HN","HK","HU","IS","IN","ID","IR","IQ","IE","IM","IL","IT","JM","JP","JE","JO","KZ","KE","KI","KP","KR","KW","KG","LA","LV","LB","LS","LR","LY","LI","LT","LU","MO","MK","MG","MW","MY","MV","ML","MT","MH","MQ","MR","MU","YT","MX","FM","MD","MC","MN","ME","MS","MA","MZ","MM","NA","NR","NP","NL","NC","NZ","NI","NE","NG","NU","NF","MP","NO","OM","PK","PW","PS","PA","PG","PY","PE","PH","PN","PL","PT","PR","QA","RE","RO","RU","RW","BL","SH","KN","LC","MF","PM","VC","WS","SM","ST","SA","SN","RS","SC","SL","SG","SX","SK","SI","SB","SO","ZA","GS","SS","ES","LK","SD","SR","SJ","SZ","SE","CH","SY","TW","TJ","TZ","TH","TL","TG","TK","TO","TT","TN","TR","TM","TC","TV","UG","UA","AE","GB","US","UM","UY","UZ","VU","VE","VN","VG","VI","WF","EH","YE","ZM","ZW")
34+
# * HTTP arg injectors
3035

31-
class _leaderboard_types(_type_check):
32-
_TYPES = ("monthly", "")
33-
34-
36+
def set_header_CSRF(): # ? if try hack me doesnt accept both at the same time, setting needs to be adjusted acordingly
37+
def predicate(ctx):
38+
ctx.add_arg(settings=["CSRF"])
39+
return arg_injector(predicate)
3540

36-
class authenticated_check:
37-
def convert(self, cls, arg):
38-
if cls.authenticated:
39-
return arg
40-
else:
41-
raise Session_required()
41+
def set_body_CSRF(): # ? if try hack me doesnt accept both at the same time, setting needs to be adjusted acordingly
42+
def predicate(ctx):
43+
ctx.add_arg(settings=["CSRF"])
44+
return arg_injector(predicate)

‎tryhackme/client.py‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,12 @@
55

66
class Client:
77
def __init__(self, session=None):
8-
self.http = HTTP(session)
8+
self.http = HTTP()
99
self._state = State(self.http)
10-
10+
11+
if(session is not None):
12+
self.login(session)
13+
1114
def login(self, session):
1215
self.http.static_login(session)
1316
if self._state.authenticated:
@@ -100,7 +103,7 @@ def server_time(self):
100103
return self.http.get_server_time().get('datetime')
101104
@property
102105
def server_stats(self):
103-
return self.http.get_server_stats()
106+
return self.http.get_site_stats()
104107
@property
105108
def subscription_cost(self):
106109
return self.http.get_subscription_cost()

‎tryhackme/cog.py‎

Lines changed: 48 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,24 @@
1+
from .errors import CheckFailed
2+
from .context import http_ctx
13

24

35
class Base_decorator:
46
pass
57

8+
class Decorator_cog:
9+
__decorator__ = Base_decorator
10+
def __init__(self, *args, **kwargs):
11+
if type(self.__decorator__) == Base_decorator:
12+
raise NotImplemented("Failed to decorate class, unknown decorator class found")
13+
14+
15+
# * Annotations auto decorator
16+
class_func_list = [i for i in type(self).__dict__ if not i.startswith('__')]
17+
for func_name in class_func_list:
18+
class_func = type(self).__dict__[func_name]
19+
setattr(self, func_name, self.__decorator__(self, class_func))
20+
super().__init__(*args, **kwargs)
21+
622
class annotater(Base_decorator):
723
def __init__(self, cls, func):
824
self.cls = cls
@@ -42,21 +58,38 @@ def convert(self, *args, **kwargs):
4258

4359
return (out_args, out_kwargs)
4460

61+
class request_annotator(annotater):
62+
def __call__(self, *args, **kwargs):
63+
try:
64+
ctx = http_ctx(state=self.cls._state, args=args, kwargs=kwargs)
65+
ctx.args, ctx.kwargs = self.convert(*ctx.args, **ctx.kwargs) # ? global context needed for convert to accept a context object
66+
self.arg_injectors(ctx)
67+
if self.check(ctx):
68+
result = self.function(self.cls, *ctx.args, **ctx.kwargs)
69+
return result
70+
raise CheckFailed()
71+
except Exception as e:
72+
print(f"Failed to run: {self.function.__name__}(), reason: {e.__repr__()}")
73+
return None
74+
75+
def arg_injectors(self, ctx):
76+
if hasattr(self.function, "__arg_injectors__"):
77+
for injector in self.function.__arg_injectors__:
78+
try:
79+
injector(ctx)
80+
except Exception as e:
81+
raise e
4582

46-
class Decorator_cog:
47-
__decorator__ = Base_decorator
48-
def __init__(self, *args, **kwargs):
49-
if type(self.__decorator__) == Base_decorator:
50-
raise NotImplemented("Failed to decorate class, unknown decorator class found")
51-
52-
53-
# * Annotations auto decorator
54-
class_func_list = [i for i in type(self).__dict__ if not i.startswith('__')]
55-
for func_name in class_func_list:
56-
class_func = type(self).__dict__[func_name]
57-
setattr(self, func_name, self.__decorator__(self, class_func))
58-
super().__init__(*args, **kwargs)
59-
83+
def check(self, ctx):
84+
if hasattr(self.function, "__checks__"):
85+
for check in self.function.__checks__:
86+
try:
87+
if not check(ctx):
88+
return False
89+
except Exception as e:
90+
raise e
91+
return True
6092

6193
class request_cog(Decorator_cog):
62-
__decorator__ = annotater
94+
__decorator__ = request_annotator
95+

‎tryhackme/context.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
from .errors import *
2+
3+
class http_ctx:
4+
def __init__(self, state, **attrs):
5+
self._state = state
6+
self.args = attrs.get("args", [])
7+
self.kwargs = attrs.get("kwargs", {})
8+
self.client_user = self._state.user
9+
10+
def add_arg(self, *args, **kwargs):
11+
if args.__len__() > 0:
12+
self.args.extend(args)
13+
elif kwargs.__len__() > 0:
14+
for key in kwargs:
15+
if key in self.kwargs:
16+
if type(self.kwargs.get(key)) is list:
17+
if type(kwargs.get(key)) is list:
18+
self.kwargs.get(key).extend(kwargs.get(key))
19+
else:
20+
self.kwargs.get(key).append(kwargs.get(key))
21+
elif type(self.kwargs.get(key)) is tuple:
22+
if type(kwargs.get(key)) is tuple:
23+
self.kwargs[key] += kwargs.get(key)
24+
elif type(kwargs.get(key)) is list:
25+
self.kwargs[key] += tuple(kwargs.get(key))
26+
else:
27+
self.kwargs[key] += (kwargs.get(key), )
28+
else:
29+
self.kwargs[key] = kwargs.get(key)
30+
else:
31+
self.kwargs[key] = kwargs.get(key)
32+
else:
33+
raise MissingArgumentError()
34+
35+
@property
36+
def authenticated(self):
37+
return self._state.authenticated

‎tryhackme/converters.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
from .errors import *
2+
3+
class _base_convert:
4+
5+
@property
6+
def name(self):
7+
return type(self).__name__.split('_')[1]
8+
9+
10+
class _type_check(_base_convert):
11+
_TYPES = ()
12+
_DEFAULT = None
13+
def convert(self, cls, arg):
14+
if arg is None: return self._DEFAULT
15+
if arg in self._TYPES: return arg
16+
else: raise TypeNotInTypeList(f"'{str(arg)}' is an invalid {self.name} type.")
17+
18+
class _not_none(_base_convert):
19+
def convert(self, cls, arg):
20+
if arg is not None:
21+
return arg
22+
else: raise NotValidUrlParameters(f"'{str(arg)}' is not allowed to be null {self.name} type.")
23+
24+
class _vpn_types(_type_check):
25+
_TYPES = ('machines', 'networks')
26+
_DEFAULT = 'machines'
27+
28+
class _county_types(_type_check):
29+
_TYPES = ("AF","AX","AL","DZ","AS","AD","AO","AI","AQ","AG","AR","AM","AW","AU","AT","AZ","BS","BH","BD","BB","BY","BE","BZ","BJ","BM","BT","BO","BQ","BA","BW","BV","BR","IO","BN","BG","BF","BI","KH","CM","CA","CV","KY","CF","TD","CL","CN","CX","CC","CO","KM","CG","CD","CK","CR","CI","HR","CU","CW","CY","CZ","DK","DJ","DM","DO","EC","EG","SV","GQ","ER","EE","ET","FK","FO","FJ","FI","FR","GF","PF","TF","GA","GM","GE","DE","GH","GI","GR","GL","GD","GP","GU","GT","GG","GN","GW","GY","HT","HM","VA","HN","HK","HU","IS","IN","ID","IR","IQ","IE","IM","IL","IT","JM","JP","JE","JO","KZ","KE","KI","KP","KR","KW","KG","LA","LV","LB","LS","LR","LY","LI","LT","LU","MO","MK","MG","MW","MY","MV","ML","MT","MH","MQ","MR","MU","YT","MX","FM","MD","MC","MN","ME","MS","MA","MZ","MM","NA","NR","NP","NL","NC","NZ","NI","NE","NG","NU","NF","MP","NO","OM","PK","PW","PS","PA","PG","PY","PE","PH","PN","PL","PT","PR","QA","RE","RO","RU","RW","BL","SH","KN","LC","MF","PM","VC","WS","SM","ST","SA","SN","RS","SC","SL","SG","SX","SK","SI","SB","SO","ZA","GS","SS","ES","LK","SD","SR","SJ","SZ","SE","CH","SY","TW","TJ","TZ","TH","TL","TG","TK","TO","TT","TN","TR","TM","TC","TV","UG","UA","AE","GB","US","UM","UY","UZ","VU","VE","VN","VG","VI","WF","EH","YE","ZM","ZW")
30+
_DEFAULT = ""
31+
32+
class _leaderboard_types(_type_check):
33+
_TYPES = ("monthly", "")
34+
_DEFAULT = ""
35+

‎tryhackme/errors.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,21 @@ class NotFound(WebError):
3232
class BaseCheckError(NotImplemented):
3333
pass
3434

35+
class CheckFailed(BaseCheckError):
36+
pass
37+
3538
class TypeNotInTypeList(BaseCheckError):
3639
pass
3740

38-
class Session_required(BaseCheckError):
41+
class SessionRequired(BaseCheckError):
42+
pass
43+
44+
45+
46+
# * Context
47+
48+
class BaseContextError(NotImplemented):
49+
pass
50+
51+
class MissingArgumentError(BaseContextError):
3952
pass

0 commit comments

Comments
 (0)