✨ Add +short flag to dig command

Author: doublevcodes

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "spade-cli"
-version = "0.2.3"
+version = "0.2.4"
 description = "You can't dig without a spade! - a super-powered networking tool!"
 authors = ["Vivaan Verma <hello@vivaanverma.com>"]
 license = "MIT"

spade/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.2.3"
+__version__ = "0.2.4"
 __app_name__ = "spade"
 
 __all__ = ("__version__", "__app_name__")

spade/commands/dig.py

@@ -1,4 +1,4 @@
-import textwrap
+from pathlib import Path
 import time
 
 import dns.rdatatype as RecordType
@@ -101,9 +101,13 @@ def dig(
     zone_file: bool = typer.Option(
         False, "-z", "--zone-file", help="Output the result in zone file format"
     ),
+    short: bool = typer.Option(False, "+short", help="Output only the relevant information"),
 ) -> None:
     """
     Run a DNS lookup.
+
+    Resolves DNS records for the given host name by querying the provided DNS
+    nameserver if one is provided, else using Cloudflare's 1.1.1.1 for speed.
     """
 
     # Perform a DNS query and get the result and the time taken to perform the query.
@@ -115,6 +119,8 @@ def dig(
         tcp=tcp,
     )
 
+    records: list[str] = lookup.rrset.to_text().split("\n")
+
     # If the user requested zone file output, we will output the result in the zone file format.
     if zone_file:
         rrset: list[str] = lookup.rrset.to_text()
@@ -123,24 +129,28 @@ def dig(
         id_: str = str(lookup.response.id)
         flags_repr: str = flags.to_text(lookup.response.flags)
 
-        output = textwrap.dedent(
-            f"""
-; <<>> Spade {__version__} <<>> {hostname_repr} 
-;; Got answer:
-;; ->>HEADER<<- opcode: {opcode_repr}, status: NOERROR, id: {id_}
-;; flags: {flags_repr};
-
-;; ANSWER SECTION:
-{rrset}
-
-;; Query time: {time_taken * 1000:.3f} msec
-;; SERVER: {lookup.nameserver}:{lookup.port}
-;; WHEN: {time.strftime("%H:%M:%S %a %d %b %Y %Z")}
-"""
+        with open(Path(__file__).parent.parent / "templates/dig/zone_file.txt", "r") as zone_file_template:
+            zone_file_template_content = zone_file_template.read()
+        
+        output = zone_file_template_content.format(
+            __version__=__version__,
+            hostname_repr=hostname_repr,
+            rrset=rrset,
+            opcode_repr=opcode_repr,
+            id_=id_,
+            flags=flags_repr,
+            ns=nameserver,
+            port=nameserver_port,
+            time_taken=time_taken,
+            timestamp=time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
         )
         typer.echo(output)
 
-    if not zone_file:
+    elif short:
+        out = "\n".join(record.split()[-1] for record in records)
+        typer.echo(out)
+
+    else:
         styled_record_type: str = typer.style(
             lookup.rdtype._name_, fg=typer.colors.BLUE
         )
@@ -153,7 +163,7 @@ def dig(
 
         output = f"\nDNS Lookup for {styled_record_type} records for {styled_hostname} through {styled_nameserver}:\n"
 
-        records: list[str] = lookup.rrset.to_text().split("\n")
+        
         styled_records: list[str] = []
 
         for record in records:

spade/commands/traceroute.py

@@ -3,13 +3,16 @@
 from icmplib.utils import resolve, unique_identifier
 from icmplib.sockets import ICMPv4Socket, ICMPv6Socket
 from icmplib.models import ICMPRequest, TimeExceeded, Hop
-from icmplib.exceptions import ICMPLibError
+from icmplib.exceptions import ICMPLibError, SocketPermissionError as SPE
 from rich.text import Text
 from rich.table import Table
 from rich.live import Live
 from ipwhois.net import Net
 from ipwhois.asn import IPASN
 from re import match
+import typer
+
+from spade.error.traceroute import SocketPermissionError
 
 def traceroute_mine(address, count=1, interval=0.0, timeout=25, first_hop=1,
         max_hops=30, id=None, source=None, family=None,
@@ -26,6 +29,10 @@ def traceroute_mine(address, count=1, interval=0.0, timeout=25, first_hop=1,
     ttl = first_hop
     host_reached = False
 
+    try:
+        _ = _Socket(source)
+    except SPE:
+        SocketPermissionError().raise_cli()
     with _Socket(source) as sock:
         while not host_reached and ttl <= max_hops:
             reply = None
@@ -82,7 +89,6 @@ def trace(host):
     table.add_column("Maximum RTT (ms)", justify="left")
     table.add_column("ASN", justify="left")
     table.add_column("ASN Info", justify="left")
-    ld = 0
     with Live(table):
         for hop in hops:
             if hop is None:
@@ -121,7 +127,6 @@ def trace(host):
                 asn,
                 f"https://bgp.tools/as/{asn}"
             )
-            ld = hop.distance
 
 def traceroute(host):
     """

spade/error/traceroute.py

@@ -0,0 +1,19 @@
+import typer
+from rich import box
+from rich import print
+from rich.panel import Panel
+
+
+class SocketPermissionError(Exception):
+    def repr(self):
+        return Panel.fit(
+            f"\n[red]This user doesn't have permission to create raw sockets on this system[/red]\n[orange]You can retry this command as an administrator or superuser[/orange]\n",
+            title=f"Missing required permissions",
+            box=box.SQUARE,
+        )
+
+    def raise_cli(self) -> str:
+        print()
+        print(self.repr())
+        print()
+        raise typer.Exit(code=1)

spade/templates/dig/zone_file.txt

@@ -0,0 +1,12 @@
+; <<>> Spade {__version__} <<>> {hostname_repr}
+;; Got answer:
+;; ->>HEADER<<- opcode: {opcode_repr}, status: NOERROR, id: {id_}
+;; flags: {flags};
+
+;; ANSWER SECTION:
+{rrset}
+
+;; Query time: {time_taken} msec
+;; SERVER: {ns}:{port}
+;; WHEN: {timestamp}
+"""