Move oracle to cipher

src/aes.zig

@@ -155,6 +155,46 @@ pub const AES = struct {
     }
 };
 
+pub const ECB_CBC_Oracle = struct {
+    mode: Mode,
+    alloc: std.mem.Allocator,
+    padded_input: []u8 = undefined,
+
+    pub fn init(alloc: std.mem.Allocator, mode: Mode) @This() {
+        return @This(){ .alloc = alloc, .mode = mode };
+    }
+
+    pub fn oracleFn(self: *@This(), alloc: std.mem.Allocator, input: []const u8) ![]u8 {
+        var rng = std.crypto.random;
+        const padding_start = rng.intRangeAtMost(u32, 0, 16);
+        const padding_end = 16 - padding_start;
+
+        self.padded_input = try self.alloc.alloc(u8, (padding_start + input.len + padding_end));
+        std.crypto.random.bytes(self.padded_input[0..padding_start]);
+        std.crypto.random.bytes(self.padded_input[(padding_start + input.len)..]);
+
+        const key = try self.alloc.alloc(u8, 16);
+        defer self.alloc.free(key);
+        std.crypto.random.bytes(key);
+
+        const iv = try self.alloc.alloc(u8, 16);
+        defer self.alloc.free(iv);
+        std.crypto.random.bytes(iv);
+
+        var cipher = switch (self.mode) {
+            .ECB => AES.init_ecb(self.alloc, key),
+            .CBC => AES.init_cbc(self.alloc, key, iv),
+        };
+        defer cipher.deinit();
+
+        return try cipher.encrypt(alloc, input);
+    }
+
+    pub fn deinit(self: *@This()) void {
+        self.alloc.free(self.padded_input);
+    }
+};
+
 test "ecb" {
     const allocator = std.testing.allocator;
 

src/main.zig

@@ -309,46 +309,6 @@ fn s2c10(allocator: std.mem.Allocator, stdout: anytype) !void {
 }
 
 fn s2c11(allocator: std.mem.Allocator, stdout: anytype) !void {
-    const Oracle = struct {
-        mode: aes.Mode,
-        alloc: std.mem.Allocator,
-        padded_input: []u8 = undefined,
-
-        fn init(alloc: std.mem.Allocator, mode: aes.Mode) @This() {
-            return @This(){ .alloc = alloc, .mode = mode };
-        }
-
-        pub fn oracleFn(self: *@This(), alloc: std.mem.Allocator, input: []const u8) ![]u8 {
-            var rng = std.crypto.random;
-            const padding_start = rng.intRangeAtMost(u32, 0, 16);
-            const padding_end = 16 - padding_start;
-
-            self.padded_input = try self.alloc.alloc(u8, (padding_start + input.len + padding_end));
-            std.crypto.random.bytes(self.padded_input[0..padding_start]);
-            std.crypto.random.bytes(self.padded_input[(padding_start + input.len)..]);
-
-            const key = try self.alloc.alloc(u8, 16);
-            defer self.alloc.free(key);
-            std.crypto.random.bytes(key);
-
-            const iv = try self.alloc.alloc(u8, 16);
-            defer self.alloc.free(iv);
-            std.crypto.random.bytes(iv);
-
-            var cipher = switch (self.mode) {
-                .ECB => aes.AES.init_ecb(self.alloc, key),
-                .CBC => aes.AES.init_cbc(self.alloc, key, iv),
-            };
-            defer cipher.deinit();
-
-            return try cipher.encrypt(alloc, input);
-        }
-
-        fn deinit(self: *@This()) void {
-            self.alloc.free(self.padded_input);
-        }
-    };
-
     for (0..1000) |_| {
         const mode_rnd = std.crypto.random.intRangeAtMost(u1, 0, 1);
         const mode = switch (mode_rnd) {
@@ -356,7 +316,7 @@ fn s2c11(allocator: std.mem.Allocator, stdout: anytype) !void {
             1 => aes.Mode.CBC,
         };
 
-        var oracle = Oracle.init(allocator, mode);
+        var oracle = aes.ECB_CBC_Oracle.init(allocator, mode);
         defer oracle.deinit();
 
         const result = try aes_crack.detect_ecb_cbc(allocator, &oracle);