use env instead of cmd args

Author: WaveringAna

README.md

@@ -1,6 +1,6 @@
 # SimpleLink
 
-A very performant and light (2MB in memory) link shortener and tracker. Written in Rust and React and uses Postgres.
+A very performant and light (2MB in memory) link shortener and tracker. Written in Rust and React and uses Postgres or SQLite.
 
 ![MainView](readme_img/mainview.jpg)
 
@@ -10,38 +10,44 @@ A very performant and light (2MB in memory) link shortener and tracker. Written
 
 ### From Docker:
 
-```Bash
+```bash
 docker run -p 8080:8080 \
     -e JWT_SECRET=change-me-in-production \
+    -e [email protected] \
+    -e SIMPLELINK_PASS=your-secure-password \
     -v simplelink_data:/data \
     ghcr.io/waveringana/simplelink:v2
 ```
 
-Find the admin-setup-token pasted into the terminal output, or in admin-setup-token.txt in the container's root.
+### Environment Variables
+
+- `JWT_SECRET`: Required. Used for JWT token generation
+- `SIMPLELINK_USER`: Optional. If set along with SIMPLELINK_PASS, creates an admin user on first run
+- `SIMPLELINK_PASS`: Optional. Admin user password
+- `DATABASE_URL`: Optional. Postgres connection string. If not set, uses SQLite
+- `INITIAL_LINKS`: Optional. Semicolon-separated list of initial links in format "url,code;url2,code2"
+- `SERVER_HOST`: Optional. Default: "127.0.0.1"
+- `SERVER_PORT`: Optional. Default: "8080"
 
-This is needed to register with the frontend. (TODO, register admin account with ENV)
+If `SIMPLELINK_USER` and `SIMPLELINK_PASS` are not passed, an admin-setup-token is pasted to the console and as a text file in the project root.
 
 ### From Docker Compose:
 
-Edit the docker-compose.yml file. It comes included with a postgressql db for use
+Edit the docker-compose.yml file. It comes included with a PostgreSQL db configuration.
 
 ## Build
 
 ### From Source
 
 First configure .env.example and save it to .env
 
-If DATABASE_URL is set, it will connect to a Postgres DB. If blank, it will use an sqlite db in /data
-
 ```bash
 git clone https://github.com/waveringana/simplelink && cd simplelink
 ./build.sh
 cargo run
 ```
 
-On an empty database, an admin-setup-token.txt is created as well as pasted into the terminal output. This is needed to make the admin account.
-
-Alternatively if you want a binary form
+Alternatively for a binary build:
 
 ```bash
 ./build.sh --binary
@@ -55,10 +61,20 @@ then check /target/release for the binary named `SimpleGit`
 docker build -t simplelink .
 docker run -p 8080:8080 \
     -e JWT_SECRET=change-me-in-production \
+    -e [email protected] \
+    -e SIMPLELINK_PASS=your-secure-password \
     -v simplelink_data:/data \
     simplelink
 ```
 
 ### From Docker Compose
 
 Adjust the included docker-compose.yml to your liking; it includes a postgres config as well.
+
+## Features
+
+- Support for both PostgreSQL and SQLite databases
+- Initial links can be configured via environment variables
+- Admin user can be created on first run via environment variables
+- Link click tracking and statistics
+- Lightweight and performant

src/main.rs

@@ -5,8 +5,11 @@ use rust_embed::RustEmbed;
 use simplelink::check_and_generate_admin_token;
 use simplelink::{create_db_pool, run_migrations};
 use simplelink::{handlers, AppState};
-use tracing::info;
+use sqlx::{Postgres, Sqlite};
+use tracing::{error, info};
 
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
 #[derive(RustEmbed)]
 #[folder = "static/"]
 struct Asset;
@@ -35,6 +38,58 @@ async fn main() -> Result<()> {
     let pool = create_db_pool().await?;
     run_migrations(&pool).await?;
 
+    // First check if admin credentials are provided in environment variables
+    let admin_credentials = match (
+        std::env::var("SIMPLELINK_USER"),
+        std::env::var("SIMPLELINK_PASS"),
+    ) {
+        (Ok(user), Ok(pass)) => Some((user, pass)),
+        _ => None,
+    };
+
+    if let Some((email, password)) = admin_credentials {
+        // Now check for existing users
+        let user_count = match &pool {
+            DatabasePool::Postgres(pool) => {
+                let mut tx = pool.begin().await?;
+                let count =
+                    sqlx::query_as::<Postgres, (i64,)>("SELECT COUNT(*)::bigint FROM users")
+                        .fetch_one(&mut *tx)
+                        .await?
+                        .0;
+                tx.commit().await?;
+                count
+            }
+            DatabasePool::Sqlite(pool) => {
+                let mut tx = pool.begin().await?;
+                let count = sqlx::query_as::<Sqlite, (i64,)>("SELECT COUNT(*) FROM users")
+                    .fetch_one(&mut *tx)
+                    .await?
+                    .0;
+                tx.commit().await?;
+                count
+            }
+        };
+
+        if user_count == 0 {
+            info!("No users found, creating admin user: {}", email);
+            match create_admin_user(&pool, &email, &password).await {
+                Ok(_) => info!("Successfully created admin user"),
+                Err(e) => {
+                    error!("Failed to create admin user: {}", e);
+                    return Err(anyhow::anyhow!("Failed to create admin user: {}", e));
+                }
+            }
+        }
+    } else {
+        info!(
+            "No admin credentials provided in environment variables, skipping admin user creation"
+        );
+    }
+
+    // Create initial links from environment variables
+    create_initial_links(&pool).await?;
+
     let admin_token = check_and_generate_admin_token(&pool).await?;
 
     let state = AppState {