Move functions over from tannerbot
This commit is contained in:
parent
3c3bf9ed37
commit
b669615f12
129
bwb/common.py
Normal file
129
bwb/common.py
Normal file
|
@ -0,0 +1,129 @@
|
|||
import base58
|
||||
import base64
|
||||
import pyotp
|
||||
import random
|
||||
import string
|
||||
|
||||
from Crypto.Cipher import AES
|
||||
from Crypto.Random import get_random_bytes
|
||||
from Crypto.Hash import SHA256, SHA512, MD5
|
||||
from Crypto.Util.Padding import pad, unpad
|
||||
|
||||
OTP_LOOKAHEAD = 25
|
||||
|
||||
class common:
|
||||
def __init__(self):
|
||||
self.otp = None
|
||||
self.secret = ''
|
||||
self.otp_count = 0
|
||||
|
||||
def init(self, secret=None):
|
||||
secret = secret or ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(12))
|
||||
encoded = secret.encode()
|
||||
b32 = base64.b32encode(encoded)
|
||||
self.otp = pyotp.HOTP(b32)
|
||||
self.secret = secret
|
||||
self.otp_count = 0
|
||||
return secret
|
||||
|
||||
def get_otp(self, target):
|
||||
if not self.otp:
|
||||
raise RuntimeError('bwb not init yet')
|
||||
if not isinstance(target, int):
|
||||
raise ValueError('invalid target')
|
||||
if int(target) <= 999999:
|
||||
raise ValueError('invalid target')
|
||||
|
||||
code = self.otp.at(self.otp_count)
|
||||
self.otp_count += 1
|
||||
return str(int(target) % int(code)).zfill(6)
|
||||
|
||||
def get_otp_cast(self):
|
||||
if not self.otp:
|
||||
raise RuntimeError('bwb not init yet')
|
||||
|
||||
code = self.otp.at(self.otp_count)
|
||||
self.otp_count += 1
|
||||
return code
|
||||
|
||||
def check_otp(self, code):
|
||||
if not self.otp:
|
||||
return False
|
||||
|
||||
for i in range(OTP_LOOKAHEAD):
|
||||
count = self.otp_count + i
|
||||
if self.otp.at(count) == str(code): # broadcast
|
||||
self.otp_count = count + 1
|
||||
return True
|
||||
elif self.TELEGRAM_ID % int(self.otp.at(count)) == int(code):
|
||||
self.otp_count = count + 1
|
||||
return True
|
||||
return False
|
||||
|
||||
def check_auth(self, text):
|
||||
try:
|
||||
int(text[:6])
|
||||
except ValueError:
|
||||
return False, None
|
||||
|
||||
if self.check_otp(text[:6]):
|
||||
return True, text[6:]
|
||||
else:
|
||||
return False, text[6:]
|
||||
|
||||
def to_b58(self, text):
|
||||
return 'l' + base58.b58encode(text.encode()).decode()
|
||||
|
||||
def from_b58(self, text):
|
||||
if not text.startswith('l'):
|
||||
return False
|
||||
try:
|
||||
return base58.b58decode(text[1:]).decode()
|
||||
except BaseException as e:
|
||||
return False
|
||||
|
||||
def enc(self, text):
|
||||
if not self.secret:
|
||||
raise RuntimeError('bwb not init yet')
|
||||
|
||||
key = SHA256.new(self.secret.encode()).digest()
|
||||
iv = get_random_bytes(4)
|
||||
cipher = AES.new(key, AES.MODE_CBC, iv * 4)
|
||||
ct = cipher.encrypt(pad(text.encode(), AES.block_size))
|
||||
return 'I' + base58.b58encode(iv + ct).decode()
|
||||
|
||||
def dec(self, ciphertext):
|
||||
if not self.secret:
|
||||
return False # so we can run every message through
|
||||
if not ciphertext.startswith('I'):
|
||||
return False
|
||||
|
||||
try:
|
||||
ciphertext = base58.b58decode(ciphertext)
|
||||
key = SHA256.new(self.secret.encode()).digest()
|
||||
iv = ciphertext[:4]
|
||||
ct = ciphertext[4:]
|
||||
cipher = AES.new(key, AES.MODE_CBC, iv * 4)
|
||||
pt = unpad(cipher.decrypt(ct), AES.block_size)
|
||||
return pt.decode()
|
||||
except BaseException as e:
|
||||
return False
|
||||
|
||||
def wrap(self, text, target=None, b58=False, enc=False):
|
||||
if target:
|
||||
code = self.get_otp(target)
|
||||
else:
|
||||
code = self.get_otp_cast()
|
||||
|
||||
text = code + text
|
||||
|
||||
if enc:
|
||||
return self.enc(text)
|
||||
elif b58:
|
||||
return self.to_b58(text)
|
||||
else:
|
||||
return text
|
||||
|
||||
def parse(self, text):
|
||||
decoded = self.dec(text) or self.from_b58(text) or text
|
||||
return self.check_auth(decoded) # returns tuple
|
6
bwb/jason.py
Normal file
6
bwb/jason.py
Normal file
|
@ -0,0 +1,6 @@
|
|||
from .common import common
|
||||
|
||||
class user(common):
|
||||
TELEGRAM_ID = 172033414
|
||||
|
||||
bwb = user()
|
6
bwb/tanner.py
Normal file
6
bwb/tanner.py
Normal file
|
@ -0,0 +1,6 @@
|
|||
from .common import common
|
||||
|
||||
class user(common):
|
||||
TELEGRAM_ID = 79316791
|
||||
|
||||
bwb = user()
|
6
bwb/tdev.py
Normal file
6
bwb/tdev.py
Normal file
|
@ -0,0 +1,6 @@
|
|||
from .common import common
|
||||
|
||||
class user(common):
|
||||
TELEGRAM_ID = 559772740
|
||||
|
||||
bwb = user()
|
Loading…
Reference in New Issue
Block a user