main.py (2679B)
1#!/usr/bin/env python3 2 3from hashlib import sha1 4from base64 import b64encode, b64decode 5from secrets import token_hex 6 7from secret import FLAG 8 9 10KEY = token_hex(16) 11 12 13def get_mac(data: bytes) -> str: 14 return sha1(KEY.encode("latin1") + data).hexdigest() 15 16 17def parse_token(token: str) -> dict: 18 # Decode token 19 token = b64decode(token) 20 21 # Check the MAC 22 token, mac = token.split(b"|mac=") 23 if get_mac(token) != mac.decode("latin1"): 24 return None 25 26 # Parse values 27 values = dict() 28 for part in token.decode("latin1").split("|"): 29 key, value = part.split("=") 30 values[key] = value 31 return values 32 33 34def generate_token(values: dict) -> str: 35 token = "|".join(f"{key}={value}" for key, value in values.items()) 36 secure_token = f"{token}|mac={get_mac(token.encode('latin1'))}" 37 38 return b64encode(secure_token.encode("latin1")).decode("latin1") 39 40 41def handle_register(): 42 name = input("What is you name? ") 43 animal = input("What is your favorite animal? ") 44 45 token = generate_token( 46 { 47 "name": name, 48 "animal": animal, 49 "admin": "false", 50 } 51 ) 52 53 print("Here is your access token:", token) 54 55 56def handle_show_animal_videos(): 57 user_data = parse_token(input("Enter access token: ")) 58 59 if user_data is None: 60 print("Invalid token.") 61 return 62 63 print( 64 f"\nHere are some {user_data['animal']} videos for you: https://www.youtube.com/results?search_query=funny+{user_data['animal']}+video+compilation" 65 ) 66 67 68def handle_show_flag(): 69 user_data = parse_token(input("Enter access token: ")) 70 71 if user_data is None: 72 print("Invalid token.") 73 return 74 75 if user_data["admin"] == "true": 76 print("The flag is", FLAG) 77 else: 78 print("You are not an admin.") 79 80 81def main(): 82 while True: 83 # Show main menu 84 85 print( 86 """ 87 1. Register 88 2. Show animal videos 89 3. Show flag 90 4. Exit 91 """ 92 ) 93 94 try: 95 choice = int(input("Enter your choice: ")) 96 except ValueError: 97 print("Please enter a number next time.") 98 continue 99 except EOFError: 100 break 101 102 if choice == 1: 103 handle_register() 104 elif choice == 2: 105 handle_show_animal_videos() 106 elif choice == 3: 107 handle_show_flag() 108 elif choice == 4: 109 break 110 else: 111 print("Please enter a valid choice.") 112 113 114if __name__ == "__main__": 115 main()