From 0fcd68c38b7fa09e1e1118ad895c60e3d183c6fe Mon Sep 17 00:00:00 2001
From: Muaz Ahmad <mahmad2000@protonmail.com>
Date: Fri, 10 May 2024 15:05:57 +0500
Subject: [PATCH] Add param to specify ip address version

---
 src/hosts.zig |  4 ++--
 src/main.zig  |  4 ++--
 src/util.zig  | 26 +++++++++++++++++++++-----
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/src/hosts.zig b/src/hosts.zig
index 4b0fd71..8349943 100644
--- a/src/hosts.zig
+++ b/src/hosts.zig
@@ -1,6 +1,6 @@
 const std = @import("std");
 const util = @import("util.zig");
 
-pub fn update_hosts(ips: util.IPs) !void {
-    std.debug.print("{any}\n", .{ips});
+pub fn update_hosts(ips: util.IPs, ip_ver: util.IP_VER_ENUM) !void {
+    std.debug.print("{any} {any}\n", .{ ips, ip_ver });
 }
diff --git a/src/main.zig b/src/main.zig
index 7fee4e7..052c082 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -5,7 +5,7 @@ const hosts = @import("hosts.zig");
 
 pub fn main() !void {
     try util.check_perms();
-    const domain = try util.get_input();
+    const domain, const ip_ver = try util.get_input();
     const ips = try mdns.get_mdns(domain);
-    try hosts.update_hosts(ips);
+    try hosts.update_hosts(ips, ip_ver);
 }
diff --git a/src/util.zig b/src/util.zig
index 33a7a74..c75f5d2 100644
--- a/src/util.zig
+++ b/src/util.zig
@@ -3,6 +3,13 @@ const std = @import("std");
 const ArgError = error{
     NotEnoughArgs,
     BadDomain,
+    InvalidAddressVer,
+};
+
+pub const IP_VER_ENUM = enum(u3) {
+    Both = 0,
+    IPv4 = 4,
+    IPv6 = 6,
 };
 
 pub const IPs = struct {
@@ -12,7 +19,7 @@ pub const IPs = struct {
 
 pub const Domain = struct {
     name: [:0]const u8,
-    labels: u8,
+    labels: u4,
 };
 
 pub fn check_perms() !void {
@@ -20,11 +27,20 @@ pub fn check_perms() !void {
     f.close();
 }
 
-pub fn get_input() !Domain {
+pub fn get_input() !struct {
+    Domain,
+    IP_VER_ENUM,
+} {
     var args = std.process.args();
-    _ = args.next();
-    const domain = args.next() orelse return ArgError.NotEnoughArgs;
-    return check_domain(domain);
+    _ = args.next(); // consume calling binary arg
+
+    const domain_str = args.next() orelse return ArgError.NotEnoughArgs;
+    const domain = try check_domain(domain_str);
+
+    const ip_ver_str = args.next() orelse return ArgError.NotEnoughArgs;
+    const ip_ver = std.meta.intToEnum(IP_VER_ENUM, try std.fmt.parseInt(u3, ip_ver_str, 10)) catch return ArgError.InvalidAddressVer;
+
+    return .{ domain, ip_ver };
 }
 
 fn check_domain(domain_str: [:0]const u8) !Domain {