xssbook/src/console.rs

use std::{net::IpAddr, collections::VecDeque, io, };
use axum::{http::{Method, Uri}, response::Response};
use lazy_static::lazy_static;
use serde::Serialize;
use serde_json::{ser::Formatter, Value};
use tokio::sync::Mutex;

use crate::types::response::ResponseCode;

struct LogMessage {
    ip: IpAddr,
    method: Method,
    uri: Uri,
    path: String,
    body: String
}

impl ToString for LogMessage {
    fn to_string(&self) -> String {
        let mut ip = self.ip.to_string();
        if ip.contains("::ffff:") {
            ip = ip.as_str()[7..].to_string()
        }
        let color = match self.method {
            Method::GET => "#3fe04f",
            Method::POST => "#853fe0",
            Method::PATCH => "#e0773f",
            Method::PUT => "#e0cb3f",
            Method::HEAD => "#3f75e0",
            Method::DELETE => "#e04c3f",
            Method::CONNECT => "#3fe0ad",
            Method::TRACE => "#e03fc5",
            Method::OPTIONS => "#423fe0",
            _ => "white"
        };
        format!("<div><span class='ip'>{}</span> <span class='method' style='color: {};'>{}</span> <span class='path'>{}{}</span> <span class='body'>{}</span></div>", ip, color, self.method, self.path, self.uri, self.body)
    }
}

lazy_static! {
    static ref LOG: Mutex<VecDeque<LogMessage>> = Mutex::new(VecDeque::with_capacity(200));
}

pub async fn log(ip: IpAddr, method: Method, uri: Uri, path: Option<String>, body: Option<String>) {

    if uri.to_string().starts_with("/console") { return; }

    let path = path.unwrap_or_default();
    let body = body.unwrap_or_default();

    tracing::info!("{} {} {}{} {}", &ip, &method, &path, &uri, &body);
    
    let message = LogMessage {
        ip, 
        method, 
        uri, 
        path, 
        body: beautify(body)
    };
    
    let mut lock = LOG.lock().await;
    if lock.len() > 200 {
        lock.pop_back();
    }
    lock.push_front(message);    
}

struct HtmlFormatter;
impl Formatter for HtmlFormatter {
    fn write_null<W>(&mut self, writer: &mut W) -> io::Result<()> where W: ?Sized + io::Write {
        writer.write_all(b"<span class='null'>null</span>")
    }

    fn write_bool<W>(&mut self, writer: &mut W, value: bool) -> io::Result<()> where W: ?Sized + io::Write {
        let s = if value {
            b"<span class='bool'>true</span>" as &[u8]
        } else {
            b"<span class='bool'>false</span>" as &[u8]
        };
        writer.write_all(s)
    }

    fn write_i8<W>(&mut self, writer: &mut W, value: i8) -> io::Result<()> where W: ?Sized + io::Write {
        let buff = format!("<span class='number'>{}</span>", value);
        writer.write_all(buff.as_bytes())
    }

    fn write_i16<W>(&mut self, writer: &mut W, value: i16) -> io::Result<()> where W: ?Sized + io::Write {
        let buff = format!("<span class='number'>{}</span>", value);
        writer.write_all(buff.as_bytes())
    }

    fn write_i32<W>(&mut self, writer: &mut W, value: i32) -> io::Result<()> where W: ?Sized + io::Write {
        let buff = format!("<span class='number'>{}</span>", value);
        writer.write_all(buff.as_bytes())
    }

    fn write_i64<W>(&mut self, writer: &mut W, value: i64) -> io::Result<()> where W: ?Sized + io::Write {
        let buff = format!("<span class='number'>{}</span>", value);
        writer.write_all(buff.as_bytes())
    }

    fn write_u8<W>(&mut self, writer: &mut W, value: u8) -> io::Result<()> where W: ?Sized + io::Write {
        let buff = format!("<span class='number'>{}</span>", value);
        writer.write_all(buff.as_bytes())
    }

    fn write_u16<W>(&mut self, writer: &mut W, value: u16) -> io::Result<()> where W: ?Sized + io::Write {
        let buff = format!("<span class='number'>{}</span>", value);
        writer.write_all(buff.as_bytes())
    }

    fn write_u32<W>(&mut self, writer: &mut W, value: u32) -> io::Result<()> where W: ?Sized + io::Write {
        let buff = format!("<span class='number'>{}</span>", value);
        writer.write_all(buff.as_bytes())
    }

    fn write_u64<W>(&mut self, writer: &mut W, value: u64) -> io::Result<()> where W: ?Sized + io::Write {
        let buff = format!("<span class='number'>{}</span>", value);
        writer.write_all(buff.as_bytes())
    }

    fn write_f32<W>(&mut self, writer: &mut W, value: f32) -> io::Result<()> where W: ?Sized + io::Write {
        let buff = format!("<span class='number'>{}</span>", value);
        writer.write_all(buff.as_bytes())
    }

    fn write_f64<W>(&mut self, writer: &mut W, value: f64) -> io::Result<()> where W: ?Sized + io::Write {
        let buff = format!("<span class='number'>{}</span>", value);
        writer.write_all(buff.as_bytes())
    }

    fn begin_string<W>(&mut self, writer: &mut W) -> io::Result<()> where W: ?Sized + io::Write {
        writer.write_all(b"<span class='string'>\"")
    }

    fn end_string<W>(&mut self, writer: &mut W) -> io::Result<()> where W: ?Sized + io::Write {
        writer.write_all(b"\"</span>")
    }

    fn begin_object_key<W>(&mut self, writer: &mut W, first: bool) -> io::Result<()> where W: ?Sized + io::Write {
        if first {
            writer.write_all(b"<span class='key'>")
        } else {
            writer.write_all(b"<span class='key'>,")
        }
    }

    fn end_object_key<W>(&mut self, writer: &mut W) -> io::Result<()> where W: ?Sized + io::Write {
        writer.write_all(b"</span>")
    }

}

fn beautify(body: String) -> String {
    if body.is_empty() {
        return "".to_string()
    }
    let Ok(mut json) = serde_json::from_str::<Value>(&body) else {
        return body
    };
    if json["password"].is_string() {
        json["password"] = Value::String("********".to_owned());
    }
    let mut writer: Vec<u8> = Vec::with_capacity(128);
    let mut serializer = serde_json::Serializer::with_formatter(&mut writer, HtmlFormatter);
    if json.serialize(&mut serializer).is_err() {
        return body
    }
    String::from_utf8_lossy(&writer).to_string()
}

pub async fn generate() -> Response {

    let lock = LOG.lock().await;

    let mut html = r#"<!DOCTYPE html>
        <html lang="en">
            <head>
                <meta charset="UTF-8">
                <meta http-equiv="refresh" content="5">
                <link rel="stylesheet" href="css/console.css">
                <title>XSSBook - Console</title>
            </head>
        <body>
        "#.to_string();

    for message in lock.iter() {
        html.push_str(&message.to_string());
    }

    html.push_str("</body></html>");

    ResponseCode::Success.html(&html)
}