derploader/src/main.zig

const std = @import("std");
const sqlite = @import("sqlite");
const clap = @import("clap");
const curl = @cImport({
    @cInclude("curl/curl.h");
});

const log = std.log.scoped(.derploader);

const params = [_]clap.Param(clap.Help){
    clap.parseParam("-h, --help          Display this help and exit.") catch unreachable,
    clap.parseParam("-c, --create=<PATH> Create new database at PATH.") catch unreachable,
    clap.parseParam("-i <ID>             Operate on ID.") catch unreachable,
    clap.parseParam("-a                  Operate on all IDs in the database.") catch unreachable,
    clap.parseParam("-m                  Download metadata for ID image.") catch unreachable,
    clap.parseParam("-d                  Download image data of ID.") catch unreachable,
    clap.parseParam("-e <PATH>           Extract image.") catch unreachable,
    clap.parseParam("-t                  Time between requests.") catch unreachable,
};

fn printFullUsage(w: anytype) !void {
    _ = try w.print("{s} ", .{std.os.argv[0]});
    try clap.usage(w, &params);
    _ = try w.writeByte('\n');
    try clap.help(w, &params);
    return;
}

fn sqliteErrorReport(str: []const u8, db: *sqlite.Db) void {
    log.err("{s}: {}", .{ str, db.getDetailedError() });
}

fn curlErrorReport(str: []const u8, code: curl.CURLcode) void {
    log.err("{s}: {s} {s}", .{ str, curl.curl_easy_strerror(code), curlerr[0.. :0] });
}

const create =
    \\CREATE TABLE IF NOT EXISTS image(
    \\  id INTEGER UNIQUE,
    \\  metadata TEXT,
    \\  image BLOB,
    \\  thumb BLOB,
    \\  full_url TEXT GENERATED ALWAYS AS
    \\    (json_extract(metadata, '$.image.representations.full')) VIRTUAL,
    \\  thumb_url TEXT GENERATED ALWAYS AS
    \\    (json_extract(metadata, '$.image.representations.thumb')) VIRTUAL,
    \\  hash_full TEXT,
    \\  hash_thumb TEXT,
    \\  hash_meta TEXT
    \\);
;

const metatable =
    \\CREATE TABLE IF NOT EXISTS derpiloader(
    \\  name TEXT,
    \\  value
    \\);
;

pub fn insertMeta(db: *sqlite.Db, id: u64, meta: []const u8) !void {
    const q =
        \\INSERT OR ROLLBACK INTO image (id, metadata) VALUES (?, ?);
    ;
    try db.exec(q, .{ .id = id, .metadata = meta });
}

const api_base = "https://derpibooru.org/api/v1/json";

var urlbuf = [_:0]u8{0} ** 512;
var curlerr = [_:0]u8{0} ** (curl.CURL_ERROR_SIZE + 1);

const hash_prefix = "blake3-";
var hash_buf = [_]u8{0} ** (std.crypto.hash.Blake3.digest_length);
var hash_buf2 = [_]u8{0} ** (std.crypto.hash.Blake3.digest_length * 2 + hash_prefix[0..].len);

fn hashit(input: []const u8) !void {
    std.crypto.hash.Blake3.hash(input, hash_buf[0..], .{});
    _ = try std.fmt.bufPrint(
        hash_buf2[0..],
        hash_prefix ++ "{s}",
        .{std.fmt.fmtSliceHexLower(hash_buf[0..])},
    );
}

pub fn main() anyerror!void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    const alloc = &gpa.allocator;

    //const key = try std.process.getEnvVarOwned(alloc, "derpikey");

    var diag = clap.Diagnostic{};
    var args = clap.parse(
        clap.Help,
        &params,
        .{ .diagnostic = &diag, .allocator = alloc },
    ) catch |err| {
        // Report useful error and exit
        diag.report(std.io.getStdErr().writer(), err) catch {};
        return;
    };
    defer args.deinit();

    if (args.flag("-h")) {
        var w = std.io.getStdOut().writer();
        try printFullUsage(w);
        return;
    }

    var db: sqlite.Db = undefined;
    const filename = "test.db3";
    try db.init(.{
        .mode = sqlite.Db.Mode{ .File = filename },
        .open_flags = .{
            .write = true,
            .create = true,
        },
        .threading_mode = .Serialized,
    });
    db.exec(create, .{}) catch sqliteErrorReport("Couldn't create table", &db);

    var ret = curl.curl_global_init(curl.CURL_GLOBAL_ALL);
    if (ret != curl.CURLE_OK) {
        log.err("cURL global init failure: {s}", .{curl.curl_easy_strerror(ret)});
        return;
    }
    defer curl.curl_global_cleanup();
    const handle = curl.curl_easy_init() orelse return error.CURLHandleInitFailed;
    defer curl.curl_easy_cleanup(handle);
    var response_buffer = std.ArrayList(u8).init(alloc);
    defer response_buffer.deinit();

    _ = curl.curl_easy_setopt(handle, curl.CURLOPT_ERRORBUFFER, &curlerr);

    const maybe_id: ?u64 = if (args.option("-i")) |id_str| blk: {
        break :blk std.fmt.parseInt(u64, id_str, 10) catch {
            log.err("Image ID must be a positive integer.", .{});
            return;
        };
    } else null;

    if (args.flag("-m")) {
        const id = if (maybe_id) |id|
            id
        else {
            log.err(
                "Operation download metadata requires an ID (-i) argument.",
                .{},
            );
            return;
        };
        const foobar = db.one(
            bool,
            "SELECT true FROM image WHERE id = ?",
            .{},
            .{ .id = id },
        ) catch {
            sqliteErrorReport("ID check read error", &db);
            return;
        };
        if (foobar) |_| {
            log.info("Info for id {d} already acquired.", .{id});
            return;
        }
        _ = try std.fmt.bufPrintZ(
            urlbuf[0..],
            api_base ++ "/images/{d}",
            .{id},
        );
        easyFetch(handle, &urlbuf, &response_buffer) catch return;
        //var w = std.io.getStdOut().writer();
        const valid = std.json.validate(response_buffer.items);

        if (valid) {
            try db.exec("BEGIN IMMEDIATE;", .{});
            errdefer db.exec("ROLLBACK;", .{}) catch {};
            insertMeta(&db, id, response_buffer.items) catch {
                sqliteErrorReport("Can't insert:", &db);
                return;
            };
            try hashit(response_buffer.items);
            db.exec(
                "UPDATE OR ROLLBACK image SET hash_meta = ? WHERE id = ?",
                .{ hash_buf2[0..], id },
            ) catch {
                sqliteErrorReport("Couldn't insert", &db);
                return;
            };
            try db.exec("COMMIT", .{});
        }
    }

    if (args.flag("-d")) {
        const id = if (maybe_id) |id|
            id
        else {
            log.err(
                "Operation download image requires an ID (-i) argument.",
                .{},
            );
            return;
        };
        const foobar = db.oneAlloc(
            struct {
                full_url: ?[:0]u8,
                thumb_url: ?[:0]u8,
            },
            alloc,
            "SELECT full_url, thumb_url FROM image WHERE id = ?",
            .{},
            .{ .id = id },
        ) catch {
            sqliteErrorReport("ID check read error", &db);
            return;
        };
        if (foobar) |res| {
            if (res.full_url) |url| {
                easyFetch(handle, url, &response_buffer) catch return;
                try db.exec("BEGIN IMMEDIATE;", .{});
                errdefer db.exec("ROLLBACK;", .{}) catch {};
                db.exec(
                    "UPDATE OR ROLLBACK image SET image = ? WHERE id = ?",
                    .{
                        .image = response_buffer.items,
                        .id = id,
                    },
                ) catch {
                    sqliteErrorReport("Couldn't add image to DB.", &db);
                    return;
                };
                try hashit(response_buffer.items);
                db.exec(
                    "UPDATE OR ROLLBACK image SET hash_full = ? WHERE id = ?",
                    .{ hash_buf2[0..], id },
                ) catch {
                    sqliteErrorReport("Couldn't add iamge hash", &db);
                    return;
                };
                try db.exec("COMMIT", .{});
                response_buffer.clearRetainingCapacity();
                std.mem.set(u8, hash_buf[0..], 0);
                std.mem.set(u8, hash_buf2[0..], 0);
            }
            if (res.thumb_url) |url| {
                easyFetch(handle, url, &response_buffer) catch return;
                try db.exec("BEGIN IMMEDIATE;", .{});
                errdefer db.exec("ROLLBACK;", .{}) catch {};
                db.exec(
                    "UPDATE OR ROLLBACK image SET thumb = ? WHERE id = ?",
                    .{
                        .thumb = response_buffer.items,
                        .id = id,
                    },
                ) catch {
                    sqliteErrorReport("Couldn't add thumb to DB", &db);
                    return;
                };
                try hashit(response_buffer.items);
                db.exec(
                    "UPDATE OR ROLLBACK image SET hash_thumb = ? WHERE id = ?",
                    .{ hash_buf2[0..], id },
                ) catch {
                    sqliteErrorReport("Couldn't add thumb hash", &db);
                    return;
                };
                try db.exec("COMMIT", .{});
            }
        } else {
            log.err("No metadata for id {d} available", .{id});
            return;
        }
    }

    if (args.option("-e")) |path| {
        const id = if (maybe_id) |id|
            id
        else {
            log.err(
                "Operation extract image requires an ID (-i) argument.",
                .{},
            );
            return;
        };
        const maybe_image = db.oneAlloc(
            []u8,
            alloc,
            "SELECT image FROM image WHERE id = ?",
            .{},
            .{ .id = id },
        ) catch {
            sqliteErrorReport("ID check read error", &db);
            return;
        };
        if (maybe_image) |image| {
            var file = try std.fs.cwd().createFile(
                path,
                .{
                    .read = false,
                    .truncate = true,
                },
            );
            try file.writeAll(image);
            file.close();
        } else {
            log.info("No image data for ID {d}.", .{id});
        }
    }
}

fn easyFetch(handle: *curl.CURL, url: [*:0]const u8, resp: *std.ArrayList(u8)) !void {
    var ret = curl.curl_easy_setopt(handle, curl.CURLOPT_URL, url);
    if (ret != curl.CURLE_OK) {
        curlErrorReport("cURL set url:", ret);
        return error.FUCK;
    }
    ret = curl.curl_easy_setopt(handle, curl.CURLOPT_WRITEFUNCTION, writeToArrayListCallback);
    if (ret != curl.CURLE_OK) {
        curlErrorReport("cURL set writefunction:", ret);
        return error.FUCK;
    }
    ret = curl.curl_easy_setopt(handle, curl.CURLOPT_WRITEDATA, resp);
    if (ret != curl.CURLE_OK) {
        curlErrorReport("cURL set writedata:", ret);
        return error.FUCK;
    }
    ret = curl.curl_easy_setopt(handle, curl.CURLOPT_USERAGENT, "Derpiloader 0.1 (linux)");
    if (ret != curl.CURLE_OK) {
        curlErrorReport("cURL set user agent:", ret);
        return error.FUCK;
    }
    ret = curl.curl_easy_perform(handle);
    if (ret != curl.CURLE_OK) {
        curlErrorReport("cURL perform:", ret);
        return error.FUCK;
    }
    log.info("Got {d} bytes", .{resp.items.len});
}

fn writeToArrayListCallback(
    data: *c_void,
    size: c_uint,
    nmemb: c_uint,
    user_data: *c_void,
) callconv(.C) c_uint {
    var buffer = @intToPtr(*std.ArrayList(u8), @ptrToInt(user_data));
    var typed_data = @intToPtr([*]u8, @ptrToInt(data));
    buffer.appendSlice(typed_data[0 .. nmemb * size]) catch return 0;
    return nmemb * size;
}