cscg24-cry2

CSCG 2024 Challenge 'Intro Crypto 2'
git clone https://git.sinitax.com/sinitax/cscg24-cry2
Log | Files | Refs | sfeed.txt

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()