const std = @import("std");
const b64 = @import("base64.zig");
const hex = @import("hex.zig");
const xor = @import("xor.zig");
const xor_crack = @import("xor_crack.zig");
const aes = @import("aes.zig");
const aes_crack = @import("aes_crack.zig");
const padding = @import("padding.zig");

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

    const args = try std.process.argsAlloc(allocator);
    defer std.process.argsFree(allocator, args);

    const stdout = std.io.getStdOut().writer();

    if (std.mem.eql(u8, args[1], "s1c1")) {
        try s1c1(allocator, stdout);
    }

    if (std.mem.eql(u8, args[1], "s1c2")) {
        try s1c2(allocator, stdout);
    }

    if (std.mem.eql(u8, args[1], "s1c3")) {
        try s1c3(allocator, stdout);
    }

    if (std.mem.eql(u8, args[1], "s1c4")) {
        try s1c4(allocator, stdout);
    }

    if (std.mem.eql(u8, args[1], "s1c5")) {
        try s1c5(allocator, stdout);
    }

    if (std.mem.eql(u8, args[1], "s1c6")) {
        try s1c6(allocator, stdout);
    }

    if (std.mem.eql(u8, args[1], "s1c7")) {
        try s1c7(allocator, stdout);
    }

    if (std.mem.eql(u8, args[1], "s1c8")) {
        try s1c8(allocator, stdout);
    }

    if (std.mem.eql(u8, args[1], "s2c09")) {
        try s2c09(allocator, stdout);
    }

    if (std.mem.eql(u8, args[1], "s2c10")) {
        try s2c10(allocator, stdout);
    }
    }
}

fn s1c1(allocator: std.mem.Allocator, stdout: anytype) !void {
    const in = "49276d206b696c6c696e6720796f757220627261696e206c696b65206120706f69736f6e6f7573206d757368726f6f6d";

    const buf = try hex.decode(allocator, in);
    defer allocator.free(buf);

    const result = try b64.encode(allocator, buf);
    defer allocator.free(result);

    try stdout.print("{s}", .{result});
}

fn s1c2(allocator: std.mem.Allocator, stdout: anytype) !void {
    const in_a = "1c0111001f010100061a024b53535009181c";
    const in_b = "686974207468652062756c6c277320657965";

    const in_a_decoded = try hex.decode(allocator, in_a);
    defer allocator.free(in_a_decoded);

    const in_b_decoded = try hex.decode(allocator, in_b);
    defer allocator.free(in_b_decoded);

    const out = try xor.xor_buffers(allocator, in_a_decoded, in_b_decoded);
    defer allocator.free(out);

    const out_encoded = try hex.encode(allocator, out);
    defer allocator.free(out_encoded);

    try stdout.print("{s}", .{out_encoded});
}

fn s1c3(allocator: std.mem.Allocator, stdout: anytype) !void {
    const in = "1b37373331363f78151b7f2b783431333d78397828372d363c78373e783a393b3736";

    const in_decoded = try hex.decode(allocator, in);
    defer allocator.free(in_decoded);

    const oracle = xor_crack.Oracle{
        .decrypt = xor.xor_byte_noalloc,
    };

    const result = try xor_crack.single(allocator, in_decoded, oracle) orelse {
        try stdout.print("No solution found", .{});
        return;
    };

    try stdout.print("Found solution key={c}, score={d}\n", .{ result[0], result[1] });

    const solution = try xor.xor_byte(allocator, in_decoded, result[0]);
    defer allocator.free(solution);

    try stdout.print("{s}\n", .{solution});
}

fn s1c4(allocator: std.mem.Allocator, stdout: anytype) !void {
    const f = @embedFile("res/4.txt");

    var f_stream = std.io.fixedBufferStream(f);
    const reader = f_stream.reader();
    const oracle = xor_crack.Oracle{
        .decrypt = xor.xor_byte_noalloc,
    };

    const out = try allocator.alloc(u8, 500);
    defer allocator.free(out);

    const threshold = 500;

    while (try reader.readUntilDelimiterOrEof(out, '\n')) |line| {
        const line_decoded = try hex.decode(allocator, line);
        defer allocator.free(line_decoded);

        const result = try xor_crack.single(allocator, line_decoded, oracle) orelse {
            continue;
        };

        if (result[1] > threshold) {
            try stdout.print("Found possible solution key={c}, score={d}\n", .{ result[0], result[1] });

            const solution = try xor.xor_byte(allocator, line_decoded, result[0]);
            defer allocator.free(solution);

            try stdout.print("{s}\n", .{solution});
        }
    }
}

fn s1c5(allocator: std.mem.Allocator, stdout: anytype) !void {
    const in =
        \\Burning 'em, if you ain't quick and nimble
        \\I go crazy when I hear a cymbal
    ;

    const key = "ICE";

    const out = try xor.xor_bytes(allocator, in, key);
    defer allocator.free(out);

    const out_encoded = try hex.encode(allocator, out);
    defer allocator.free(out_encoded);

    try stdout.print("{s}", .{out_encoded});
}

fn s1c6(allocator: std.mem.Allocator, stdout: anytype) !void {
    // prepare input
    const f = @embedFile("res/6.txt");

    var f_stream = std.io.fixedBufferStream(f);
    const reader = f_stream.reader();

    var joined_lines = std.ArrayList(u8).init(allocator);
    defer joined_lines.deinit();
    const joined_lines_writer = joined_lines.writer();

    while (true) {
        reader.streamUntilDelimiter(joined_lines_writer, '\n', null) catch |err| {
            if (err == error.EndOfStream) break else return err;
        };
    }

    const input_decoded = try b64.decode(allocator, joined_lines.items);
    defer allocator.free(input_decoded);

    // crack key
    const oracle = xor_crack.Oracle{
        .decrypt = xor.xor_byte_noalloc,
    };
    const key = try xor_crack.multi(allocator, input_decoded, oracle) orelse {
        try stdout.print("Could not find solution", .{});
        return;
    };
    defer allocator.free(key);

    // print result
    const result = try xor.xor_bytes(allocator, input_decoded, key);
    defer allocator.free(result);

    try stdout.print("{s}", .{result});
}

fn s1c7(allocator: std.mem.Allocator, stdout: anytype) !void {
    // prepare input
    const f = @embedFile("res/7.txt");

    var f_stream = std.io.fixedBufferStream(f);
    const reader = f_stream.reader();

    var joined_lines = std.ArrayList(u8).init(allocator);
    defer joined_lines.deinit();
    const joined_lines_writer = joined_lines.writer();

    while (true) {
        reader.streamUntilDelimiter(joined_lines_writer, '\n', null) catch |err| {
            if (err == error.EndOfStream) break else return err;
        };
    }

    const input_decoded = try b64.decode(allocator, joined_lines.items);
    defer allocator.free(input_decoded);

    // decrypt
    const key = "YELLOW SUBMARINE";

    var cipher = aes.AES.init(allocator, key);
    defer cipher.deinit();

    const result = try cipher.decrypt(allocator, input_decoded);
    defer allocator.free(result);

    try stdout.print("{s}", .{result});
}

fn s1c8(allocator: std.mem.Allocator, stdout: anytype) !void {
    // prepare input
    const f = @embedFile("res/8.txt");

    var input_stream = std.io.fixedBufferStream(f);
    const reader = input_stream.reader();

    var output_buf = std.ArrayList(u8).init(allocator);
    defer output_buf.deinit();

    const writer = output_buf.writer();

    var lineno: usize = 1;
    while (true) {
        defer output_buf.clearRetainingCapacity();

        reader.streamUntilDelimiter(writer, '\n', null) catch |err| {
            switch (err) {
                error.EndOfStream => break,
                else => return err,
            }
        };

        const is_ecb: bool = try aes_crack.detect(allocator, output_buf.items);
        if (is_ecb) {
            try stdout.print("Found vulnerable cipher lineno: {d}\n{s}\n", .{ lineno, output_buf.items });
        }

        lineno += 1;
    }
}

fn s2c09(allocator: std.mem.Allocator, stdout: anytype) !void {
    const input = "YELLOW SUBMARINE";

    const result = try padding.pkcs7(allocator, input, 20);
    defer allocator.free(result);

    try stdout.print("{s}", .{result});
}

fn s2c10(allocator: std.mem.Allocator, stdout: anytype) !void {
    // prepare input
    const f = @embedFile("res/10.txt");

    var f_stream = std.io.fixedBufferStream(f);
    const reader = f_stream.reader();

    var joined_lines = std.ArrayList(u8).init(allocator);
    defer joined_lines.deinit();
    const joined_lines_writer = joined_lines.writer();

    while (true) {
        reader.streamUntilDelimiter(joined_lines_writer, '\n', null) catch |err| {
            if (err == error.EndOfStream) break else return err;
        };
    }

    const input_decoded = try b64.decode(allocator, joined_lines.items);
    defer allocator.free(input_decoded);

    // decrypt
    const key = "YELLOW SUBMARINE";
    const iv = [_]u8{0x00} ** 16;

    var cipher = aes.AES.init_cbc(allocator, key, &iv);
    defer cipher.deinit();

    const result = try cipher.decrypt(allocator, input_decoded);
    defer allocator.free(result);

    try stdout.print("{s}", .{result});
}