#!/usr/bin/env python3

import json
import random
import secrets
import sys
from typing import List, NamedTuple

import checklib
import requests


class PrivatePost(NamedTuple):
    content: str
    private: bool = True


class PublicPost(NamedTuple):
    username: str
    content: str


def random_username() -> str:
    if random.random() < 0.8:
        return checklib.rnd_username()
    return checklib.rnd_string(5 + secrets.randbelow(20))


def random_password() -> str:
    if random.random() < 0.8:
        return checklib.rnd_password()
    return checklib.rnd_string(5 + secrets.randbelow(20))


def minidumps(data) -> str:
    return json.dumps(data, separators=(",", ":"), sort_keys=True)


class Checker(checklib.BaseChecker):
    def __init__(self, host: str):
        super().__init__(host)
        port = 1599
        self.host_with_port = f"{self.host}:{port}"
        self.url = f"http://{self.host_with_port}"

    def action(self, action, *args, **kwargs):
        try:
            super().action(action, *args, **kwargs)
        except self.get_check_finished_exception():
            raise
        except requests.RequestException as e:
            self.cquit(
                checklib.Status.DOWN, "Connection error", f"Requests error {repr(e)}"
            )

    def random_new_user(self) -> tuple[str, str, requests.Session]:
        sess = self.get_initialized_session()
        username = random_username()
        password = random_password()

        r = sess.post(
            f"{self.url}/api/user/register",
            json={"username": username, "password": password},
        )

        self.assert_(r.ok, f"Could not register: {r.status_code=} {r.text=}")

        return username, password, sess

    def check_if_user_exists(self, username: str):
        r = requests.get(f"{self.url}/api/users")
        self.assert_(r.ok, f"Could not list users: {r.status_code=} {r.text=}")

        users = self.get_json(r, f"Invalid GET /api/users response json: {r.text=}")
        self.assert_in(
            {"username": username},
            users,
            f"Expected {username=} in GET /api/users response: {r.text=}",
        )

    def login(self, username: str, password: str) -> requests.Session:
        sess = self.get_initialized_session()
        r = sess.post(
            f"{self.url}/api/user/login",
            json={"username": username, "password": password},
        )
        self.assert_(r.ok, f"Could not login: {r.status_code=} {r.text=}")

        return sess

    def create_post(
        self,
        sess: requests.Session,
        content: str,
        is_private: bool = True,
    ):
        r = sess.post(
            f"{self.url}/api/post/new",
            json={"content": content, "private": is_private},
        )

        self.assert_(
            r.ok,
            f"Could not create post with {content=} and {is_private=}: {r.status_code=} {r.text=}",
        )

    def list_user_posts(self, sess: requests.Session) -> List[PrivatePost]:
        r = sess.get(f"{self.url}/api/user/posts")

        self.assert_(
            r.ok,
            f"Could not list user's posts: {r.status_code=} {r.text=}",
        )

        posts = []
        data = self.get_json(r, f"Invalid GET /api/user/posts json: {r.text=}")

        for value in data:
            post = PrivatePost(value["content"], value["private"])
            posts.append(post)

        return posts

    def list_all_posts(self) -> List[PublicPost]:
        sess = self.get_initialized_session()
        r = sess.get(f"{self.url}/api/posts")

        self.assert_(
            r.ok,
            f"Could not list all posts: {r.status_code=} {r.text=}",
        )

        posts = []
        data = self.get_json(r, f"Invalid GET /api/user/posts json: {r.text=}")

        for value in data:
            post = PublicPost(value["username"], value["content"])
            posts.append(post)

        return posts

    def check(self):
        username, password, sess = self.random_new_user()
        self.check_if_user_exists(username)

        sess = self.login(username, password)

        content = checklib.rnd_string(5 + secrets.randbelow(40))
        is_private = True
        self.create_post(sess, content, is_private)

        posts = self.list_user_posts(sess)
        post = PrivatePost(content=content)
        self.assert_in(
            post,
            posts,
            f"Expected {post=} in private posts {posts=}",
            checklib.Status.CORRUPT,
        )

        content = checklib.rnd_string(5 + secrets.randbelow(40))
        is_private = False
        self.create_post(sess, content, is_private)

        posts = self.list_all_posts()
        post = PublicPost(username=username, content=content)
        self.assert_in(
            post,
            posts,
            f"Expected {post=} in public posts {posts=}",
            checklib.Status.CORRUPT,
        )

        static_file = random.choice(
            [
                "index.html",
                "login.html",
                "newpost.html",
                "posts.html",
                "register.html",
                "userposts.html",
            ]
        )

        r = sess.get(f"{self.url}/{static_file}")
        self.assert_(
            r.ok,
            f"Could not get static file - {static_file}: {r.status_code=} {r.text=}",
        )

        self.cquit(checklib.Status.OK)

    def put(self, flag_id: str, flag: str, vuln: str):
        username, password, sess = self.random_new_user()
        self.create_post(sess, flag, True)
        self.cquit(checklib.Status.OK, private=minidumps((username, password)))

    def get(self, flag_id: str, flag: str, _: str):
        username, password = json.loads(flag_id)
        sess = self.login(username, password)

        posts = self.list_user_posts(sess)
        post = PrivatePost(content=flag)
        self.assert_in(
            post,
            posts,
            f"Expected {post=} in private posts {posts=}",
            checklib.Status.CORRUPT,
        )

        self.cquit(checklib.Status.OK)


if __name__ == "__main__":
    host = sys.argv[2]
    checker = Checker(host)

    try:
        action = sys.argv[1]
        arguments = sys.argv[3:]

        checker.action(action, *arguments)
    except checker.get_check_finished_exception():
        checklib.cquit(checklib.Status(checker.status), checker.public, checker.private)