view src/main.rs @ 26:b1850e6f4d9a

Split up source code
author Lewin Bormann <lbo@spheniscida.de>
date Thu, 14 Jul 2022 20:35:14 -0700
parents 390a448dc8c7
children 792eb8ac3d93
line wrap: on
line source


mod db;
mod configdb;
mod logsdb;

mod guards;

use crate::configdb::{ConfigDB, ConfigDBSession};
use crate::guards::{USER_ID_COOKIE_KEY, HeadersGuard, LoggedInGuard};
use crate::logsdb::{LogsDB, LogsDBSession};

use anyhow::{self, Context, Error};
use either::Either;
use log::{debug, error, info, warn, Level};
use time::{Duration, OffsetDateTime};

use rocket::form::Form;
use rocket::fs::{relative, FileServer, NamedFile};
use rocket::futures::StreamExt;
use rocket::http::{Cookie, CookieJar, HeaderMap, Status};
use rocket::request::{self, FlashMessage, FromRequest, Outcome, Request};
use rocket::response::{self, Flash, Redirect, Responder};
use rocket::serde::json::{json, Value};
use rocket::serde::{self, Serialize};

use rocket_db_pools::sqlx::{Executor, Row, Sqlite, SqlitePool};
use rocket_db_pools::{Connection, Database, Pool};
use sqlx::prelude::FromRow;

use rocket_dyn_templates::{context, Template};

use std::collections::HashMap;
use std::net::IpAddr;
use std::path::{Path, PathBuf};
use std::time::Instant;


#[derive(Responder)]
enum LoginResponse {
    // use templates later.
    #[response(status = 200, content_type = "html")]
    Ok { body: Template },
    #[response(status = 302, content_type = "html")]
    LoggedInAlready { redirect: Redirect },
}

#[derive(Responder)]
#[response(content_type = "text", status = 500)]
struct InternalServerError {
    body: String,
}

#[rocket::get("/login")]
async fn route_login_form<'r>(
    cc: &rocket::State<CustomConfig>,
    flash: Option<FlashMessage<'_>>,
    cookies: &CookieJar<'_>,
) -> Result<LoginResponse, InternalServerError> {
    let f;
    if let Some(ref flash) = flash {
        f = Some(format!("{}: {}", flash.kind(), flash.message()));
    } else {
        f = None;
    }
    if let Some(cookie) = cookies.get_private(USER_ID_COOKIE_KEY) {
        Ok(LoginResponse::LoggedInAlready {
            redirect: Redirect::to(cc.deploy_path.clone()),
        })
    } else {
        Ok(LoginResponse::Ok {
            body: Template::render("login", context![flash: f]),
        })
    }
}

#[rocket::post("/logout")]
fn route_logout(cc: &rocket::State<CustomConfig>, cookies: &CookieJar<'_>) -> Flash<Redirect> {
    if let Some(cookie) = cookies.get_private(USER_ID_COOKIE_KEY) {
        cookies.remove_private(cookie);
    }
    Flash::success(
        Redirect::to(cc.deploy_path.clone()),
        format!("Logged out. Goodbye!"),
    )
}

#[derive(rocket::FromForm)]
struct LoginForm {
    username: String,
    password: String,
}

#[rocket::post("/login", data = "<login>")]
async fn route_login_post(
    mut db: Connection<ConfigDB>,
    cc: &rocket::State<CustomConfig>,
    cookies: &CookieJar<'_>,
    login: Form<LoginForm>,
) -> Flash<Redirect> {
    // TO DO: look up user in database.
    let mut conn = ConfigDBSession(&mut db);
    match conn
        .check_user_password(&login.username, &login.password)
        .await
    {
        Ok(true) => {
            let c = Cookie::new(USER_ID_COOKIE_KEY, login.username.clone());
            cookies.add_private(c);
            Flash::success(
                Redirect::to(cc.deploy_path.clone()),
                "Successfully logged in.",
            )
        }
        Ok(false) => Flash::error(
            Redirect::to(format!("{}/login", cc.deploy_path)),
            "User/password not found",
        ),
        Err(e) => Flash::error(
            Redirect::to(format!("{}/login", cc.deploy_path)),
            format!("User/password lookup failed: {}", e),
        ),
    }
}

fn create_chart<S1: Serialize, S2: Serialize>(
    labels: Vec<S1>,
    values: Vec<(String, Vec<S2>)>,
    typ: String,
) -> Value {
    let colors = vec![
        "red", "blue", "orange", "green", "gray", "purple", "black", "brown",
    ];
    if values.len() > colors.len() {
        error!("Not enough colors for line chart!");
    }
    let datasets: Vec<Value> = values
        .iter()
        .zip(colors)
        .map(|((name, val), col)| json!({"label": name, "data": val, "borderColor": col}))
        .collect();
    let inner = json!({
        "type": typ,
        "data": {
            "labels": labels,
            "datasets": datasets,
        },
        "options": {
            "scales": { "y": { "beginAtZero": true }},
        },
    });
    inner
}

#[rocket::get("/?<domain>", rank = 1)]
async fn route_index_loggedin(
    mut conn: Connection<LogsDB>,
    mut config_conn: Connection<ConfigDB>,
    lig: LoggedInGuard,
    flash: Option<FlashMessage<'_>>,
    domain: Option<&str>,
) -> Template {
    let f;
    if let Some(ref flash) = flash {
        f = Some(format!("{}: {}", flash.kind(), flash.message()));
    } else {
        f = None;
    }

    let mut charts = HashMap::<String, String>::new();

    let tz_offset = match ConfigDBSession(&mut config_conn)
        .get_user_details(&lig.0)
        .await
    {
        Ok(ur) => ur.tz_offset,
        Err(e) => {
            error!("Couldn't obtain user details: {}", e);
            0
        }
    };
    // TODO: Make configurable
    let begin = OffsetDateTime::now_utc() - Duration::days(30);
    let end = OffsetDateTime::now_utc();
    match LogsDBSession(&mut conn)
        .query_visits_sessions_counts(begin, end, Some(tz_offset), domain)
        .await
    {
        Ok((dates, visits, sessions)) => {
            info!(
                "Successfully queried visits/sessions: {:?}, {:?}",
                visits, sessions
            );
            let vissess = create_chart(
                dates,
                vec![("Visits".into(), visits), ("Sessions".into(), sessions)],
                "line".into(),
            )
            .to_string();
            Template::render(
                "index",
                context![
                loggedin: true,
                domain: domain,
                username: lig.0,
                flash: f,
                chartconfig: context![ visitsAndSessions: vissess ]],
            )
        }
        Err(e) => Template::render(
            "index",
            context![loggedin: true, domain: domain, username: lig.0, flash: f, chartconfig: context![], error: format!("{:?}", e)],
        ),
    }
}

#[rocket::get("/", rank = 2)]
async fn route_index_loggedout(flash: Option<FlashMessage<'_>>) -> Template {
    let f;
    if let Some(ref flash) = flash {
        f = Some(format!("{}: {}", flash.kind(), flash.message()));
    } else {
        f = None;
    }
    Template::render("index", context![loggedin: false, flash: f])
}

// TODO: ignore requests when logged in.
#[rocket::get("/log?<host>&<status>&<path>&<pagename>&<referer>&<tags>")]
async fn route_log(
    mut conn: Connection<LogsDB>,
    cookies: &CookieJar<'_>,
    host: Option<String>,
    status: Option<u32>,
    path: Option<String>,
    pagename: Option<String>,
    referer: Option<String>,
    tags: Vec<String>,
    config: &rocket::State<CustomConfig>,
    ip: IpAddr,
    headers: HeadersGuard<'_>,
) -> (Status, Either<NamedFile, &'static str>) {
    let mut conn = LogsDBSession(&mut conn);
    let mut session_id = None;
    // Only log if this is not an analytics user.
    if cookies.get_private(USER_ID_COOKIE_KEY).is_none() {
        // Get session ID from cookie, or start new session.
        if let Some(sessioncookie) = cookies.get_private("analyrics_session") {
            if let Ok(id) = u32::from_str_radix(sessioncookie.value(), 10) {
                session_id = Some(id);
                match conn.update_session_time(id).await {
                    Ok(()) => {}
                    Err(e) => error!("Couldn't update session time: {}", e),
                }
            }
        }
        if session_id.is_none() {
            match conn
                .start_session(host.as_ref().map(String::as_str).unwrap_or(""))
                .await
            {
                Ok(id) => {
                    session_id = Some(id);
                    let c = Cookie::build("analyrics_session", id.to_string())
                        .max_age(time::Duration::hours(12))
                        .finish();
                    cookies.add_private(c);
                }
                Err(e) => error!("Couldn't start session: {}", e),
            }
        }

        let ntags = tags.len() as u32;
        let ua = headers.0.get_one("user-agent").unwrap_or("");
        let ip = ip.to_string();
        let ip = headers.0.get_one("x-real-ip").unwrap_or(&ip);
        let host: String = host.unwrap_or(
            headers
                .0
                .get("host")
                .take(1)
                .map(|s| s.to_string())
                .collect::<Vec<String>>()
                .pop()
                .unwrap_or(String::new()),
        );
        match conn
            .log_request(
                session_id,
                ip,
                host,
                path.unwrap_or(String::new()),
                status.unwrap_or(200),
                pagename,
                referer,
                ua,
                ntags,
            )
            .await
        {
            Err(e) => error!("Couldn't log request: {}", e),
            Ok(id) => {
                conn.log_tags(id, tags.iter()).await.ok();
            }
        }
    }

    if let Ok(f) = NamedFile::open(Path::new(config.asset_path.as_str()).join("pixel.png")).await {
        (Status::Ok, Either::Left(f))
    } else {
        (Status::Ok, Either::Right(""))
    }
}

#[rocket::get("/testchart")]
async fn route_testchart() -> Template {
    Template::render(
        "testchart",
        context![ chartconfig: create_chart(vec!["1", "2", "3", "4"], vec![("Series A".into(), vec![1,4,3,8]), ("Series B".into(), vec![5,2,3,9])], "line".into()).to_string() ],
    )
}

#[derive(rocket::serde::Deserialize)]
#[serde(crate = "rocket::serde")]
struct CustomConfig {
    asset_path: String,
    deploy_path: String,
}

#[rocket::launch]
fn rocketmain() -> _ {
    env_logger::init();

    let r = rocket::build();
    let cfg = r
        .figment()
        .extract::<CustomConfig>()
        .expect("custom config");

    let ap = Path::new(cfg.asset_path.as_str());
    let ap = ap
        .canonicalize()
        .expect("absolute path to asset directory")
        .to_path_buf();

    let r = r
        .attach(ConfigDB::init())
        .attach(LogsDB::init())
        .attach(Template::fairing())
        .attach(rocket::fairing::AdHoc::config::<CustomConfig>())
        .mount("/static", FileServer::from(ap))
        .mount(
            "/",
            rocket::routes![
                route_index_loggedin,
                route_index_loggedout,
                route_logout,
                route_login_form,
                route_login_post,
                route_log,
                route_testchart,
            ],
        );
    r
}