|
| 1 | +import math |
| 2 | +import random |
| 3 | +import re |
| 4 | +import mmh3 |
| 5 | + |
| 6 | +def randbytes(n): return bytes ([random.randint(0,255) for i in range(n)]) |
| 7 | + |
| 8 | +class BloomFilter: |
| 9 | + def __init__(self, m, k, hash_func=mmh3.hash): |
| 10 | + self.__m = m |
| 11 | + self.__k = k |
| 12 | + self.__i = 0 |
| 13 | + self.__digests = set() |
| 14 | + self.hash = hash_func |
| 15 | + |
| 16 | + def security(self): |
| 17 | + false_positive = pow( |
| 18 | + 1 - pow(math.e, -self.__k * self.__i / self.__m), self.__k) |
| 19 | + try: |
| 20 | + return int(1 / false_positive).bit_length() |
| 21 | + except (ZeroDivisionError, OverflowError): |
| 22 | + return float('inf') |
| 23 | + |
| 24 | + def _add(self, item): |
| 25 | + self.__i += 1 |
| 26 | + for i in range(self.__k): |
| 27 | + self.__digests.add(self.hash(item, i) % self.__m) |
| 28 | + |
| 29 | + def check(self, item): |
| 30 | + return all(self.hash(item, i) % self.__m in self.__digests |
| 31 | + for i in range(self.__k)) |
| 32 | + |
| 33 | + def num_passwords(self): |
| 34 | + return self.__i |
| 35 | + |
| 36 | + def memory_consumption(self): |
| 37 | + return 4*len(self.__digests) |
| 38 | + |
| 39 | + |
| 40 | +class PasswordDB(BloomFilter): |
| 41 | + def __init__(self, m, k, security, hash_func=mmh3.hash): |
| 42 | + super().__init__(m, k, hash_func) |
| 43 | + self.add_keys(security) |
| 44 | + self.addition_quota = 1 |
| 45 | + self.added_keys = set() |
| 46 | + |
| 47 | + def add_keys(self, thresh_security): |
| 48 | + while self.security() > thresh_security: |
| 49 | + self._add(randbytes(256)) |
| 50 | + print("Added {} security keys to DB".format(self.num_passwords())) |
| 51 | + print("Original size of keys {} KB vs {} KB in DB".format( |
| 52 | + self.num_passwords()//4, self.memory_consumption()//1024)) |
| 53 | + |
| 54 | + def check_admin(self, key): |
| 55 | + if not re.match(b".{32,}", key): |
| 56 | + print("Admin key should be atleast 32 characters long") |
| 57 | + return False |
| 58 | + if not re.match(b"(?=.*[a-z])", key): |
| 59 | + print("Admin key should contain atleast 1 lowercase character") |
| 60 | + return False |
| 61 | + if not re.match(b"(?=.*[A-Z])", key): |
| 62 | + print("Admin key should contain atleast 1 uppercase character") |
| 63 | + return False |
| 64 | + if not re.match(br"(?=.*\d)", key): |
| 65 | + print("Admin key should contain atleast 1 digit character") |
| 66 | + return False |
| 67 | + if not re.match(br"(?=.*\W)", key): |
| 68 | + print("Admin key should contain atleast 1 special character") |
| 69 | + return False |
| 70 | + if key in self.added_keys: |
| 71 | + print("Admin account restricted for free tier") |
| 72 | + return False |
| 73 | + return self.check(key) |
| 74 | + |
| 75 | + def query_db(self, key): |
| 76 | + if self.check(key): |
| 77 | + print("Key present in DB") |
| 78 | + else: |
| 79 | + print("Key not present in DB") |
| 80 | + |
| 81 | + def add_sample(self, key): |
| 82 | + if self.addition_quota > 0: |
| 83 | + self._add(key) |
| 84 | + self.added_keys.add(key) |
| 85 | + self.addition_quota -= 1 |
| 86 | + print("key added successfully to DB") |
| 87 | + else: |
| 88 | + print("API quota exceeded") |
| 89 | + |
| 90 | + |
| 91 | +BANNER = r""" |
| 92 | + ____ ____ ____ ____ ____ ___ ____ ____ _ _ ____ |
| 93 | +( _ \(_ _)( ___)( ___)( ___)/ __)(_ _)( ___)( \( )(_ _) |
| 94 | + )(_) )_)(_ )__) )__) )__)( (__ _)(_ )__) ) ( )( |
| 95 | +(____/(____)(__) (__) (____)\___)(____)(____)(_)\_) (__) |
| 96 | +
|
| 97 | +Welcome to diffecient security key database API for securely |
| 98 | +and efficiently saving tonnes of long security keys! |
| 99 | +Feel FREE to query your security keys and pay a little to |
| 100 | +add your own security keys to our state of the art DB! |
| 101 | +We trust our product so much that we even save our own keys here |
| 102 | +""" |
| 103 | +print(BANNER) |
| 104 | +PASSWORD_DB = PasswordDB(2**32 - 5, 47, 768, mmh3.hash) |
| 105 | +while True: |
| 106 | + try: |
| 107 | + option = int(input("Enter API option:\n")) |
| 108 | + if option == 1: |
| 109 | + key = bytes.fromhex(input("Enter key in hex\n")) |
| 110 | + PASSWORD_DB.query_db(key) |
| 111 | + elif option == 2: |
| 112 | + key = bytes.fromhex(input("Enter key in hex\n")) |
| 113 | + PASSWORD_DB.add_sample(key) |
| 114 | + elif option == 3: |
| 115 | + key = bytes.fromhex(input("Enter key in hex\n")) |
| 116 | + if PASSWORD_DB.check_admin(key): |
| 117 | + from flag import flag |
| 118 | + print(flag) |
| 119 | + else: |
| 120 | + print("No Admin no flag") |
| 121 | + elif option == 4: |
| 122 | + exit(0) |
| 123 | + except: |
| 124 | + print("Something wrong happened") |
| 125 | + exit(1) |
0 commit comments