1 files changed, 115 insertions, 0 deletions
diff --git a/solve/main.py b/solve/main.py
new file mode 100755
index 0000000..e3f4134
--- /dev/null
+++ b/solve/main.py
@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+
+from hashlib import sha1
+from base64 import b64encode, b64decode
+from secrets import token_hex
+
+from secret import FLAG
+
+
+KEY = token_hex(16)
+
+
+def get_mac(data: bytes) -> str:
+    return sha1(KEY.encode("latin1") + data).hexdigest()
+
+
+def parse_token(token: str) -> dict:
+    # Decode token
+    token = b64decode(token)
+
+    # Check the MAC
+    token, mac = token.split(b"|mac=")
+    if get_mac(token) != mac.decode("latin1"):
+        return None
+
+    # Parse values
+    values = dict()
+    for part in token.decode("latin1").split("|"):
+        key, value = part.split("=")
+        values[key] = value
+    return values
+
+
+def generate_token(values: dict) -> str:
+    token = "|".join(f"{key}={value}" for key, value in values.items())
+    secure_token = f"{token}|mac={get_mac(token.encode('latin1'))}"
+
+    return b64encode(secure_token.encode("latin1")).decode("latin1")
+
+
+def handle_register():
+    name = input("What is you name? ")
+    animal = input("What is your favorite animal? ")
+
+    token = generate_token(
+        {
+            "name": name,
+            "animal": animal,
+            "admin": "false",
+        }
+    )
+
+    print("Here is your access token:", token)
+
+
+def handle_show_animal_videos():
+    user_data = parse_token(input("Enter access token: "))
+
+    if user_data is None:
+        print("Invalid token.")
+        return
+
+    print(
+        f"\nHere are some {user_data['animal']} videos for you: https://www.youtube.com/results?search_query=funny+{user_data['animal']}+video+compilation"
+    )
+
+
+def handle_show_flag():
+    user_data = parse_token(input("Enter access token: "))
+
+    if user_data is None:
+        print("Invalid token.")
+        return
+
+    if user_data["admin"] == "true":
+        print("The flag is", FLAG)
+    else:
+        print("You are not an admin.")
+
+
+def main():
+    while True:
+        # Show main menu
+
+        print(
+            """
+        1. Register
+        2. Show animal videos
+        3. Show flag
+        4. Exit
+        """
+        )
+
+        try:
+            choice = int(input("Enter your choice: "))
+        except ValueError:
+            print("Please enter a number next time.")
+            continue
+        except EOFError:
+            break
+
+        if choice == 1:
+            handle_register()
+        elif choice == 2:
+            handle_show_animal_videos()
+        elif choice == 3:
+            handle_show_flag()
+        elif choice == 4:
+            break
+        else:
+            print("Please enter a valid choice.")
+
+
+if __name__ == "__main__":
+    main()