From afecb475a5f311ba1b0f97a74a8ba09d547ef8c1 Mon Sep 17 00:00:00 2001
From: guanana <guanana2@gmail.com>
Date: Tue, 15 Dec 2020 01:27:34 +0000
Subject: [PATCH] __init__.py: Improve tag checking (With the previous version
 of check wasn't working, don't know if it was dependant on old Netbox
 version) nmap.py: Add extra function to improve DNS resolution since nmap is
 not always consistent

---
 nbs/__init__.py | 11 ++++++++---
 nbs/nmap.py     | 27 +++++++++++++++++++++------
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/nbs/__init__.py b/nbs/__init__.py
index ee79cc3..a3ee322 100644
--- a/nbs/__init__.py
+++ b/nbs/__init__.py
@@ -1,6 +1,6 @@
 import logging
-import requests
 
+import requests
 from pynetbox import api
 
 
@@ -47,7 +47,12 @@ class NetBoxScanner(object):
             return False
 
         if nbhost:
-            if (self.tag in nbhost.tags):
+            tag_update = False
+            for tag in nbhost.tags:
+                if (self.tag == tag.name):
+                    tag_update = True
+                    break
+            if tag_update:
                 if (host[1] != nbhost.description):
                     aux = nbhost.description
                     nbhost.description = host[1]
@@ -59,7 +64,7 @@ class NetBoxScanner(object):
                     logging.info(f'unchanged: {host[0]}/32 "{host[1]}"')
                     self.stats['unchanged'] += 1
             else:
-                logging.info(f'unchanged: {host[0]}/32 "{host[1]}"')
+                logging.info(f'no-tag(unchanged): {host[0]}/32 "{host[1]}"')
                 self.stats['unchanged'] += 1
         else:
             self.netbox.ipam.ip_addresses.create(
diff --git a/nbs/nmap.py b/nbs/nmap.py
index 168a77b..2850530 100644
--- a/nbs/nmap.py
+++ b/nbs/nmap.py
@@ -1,5 +1,8 @@
+import socket
+
 import nmap3
 
+
 class Nmap(object):
 
     def __init__(self, unknown, networks):
@@ -9,17 +12,30 @@ class Nmap(object):
         self.scan_results = {}
 
     def scan(self):
-        nmap = nmap3.NmapHostDiscovery() # instantiate nmap object
+        nmap = nmap3.NmapHostDiscovery()  # instantiate nmap object
         for item in self.networks:
             temp_scan_result = nmap.nmap_no_portscan(item.replace('\n', ''))
             self.scan_results = {**self.scan_results, **temp_scan_result}
+            self.scan_results.pop("stats")
+            self.scan_results.pop("runtime")
         return self.scan_results
 
+    def dns_resolution(self):
+        # Try to improve DNS resolution since NMAP is not consistent
+        for ip, v in self.scan_results.items():
+            try:
+                name, arpa, ip = socket.gethostbyaddr(ip)
+                try:
+                    v["hostname"][0]["name"]
+                except (TypeError, IndexError):
+                    v.update({"hostname": {"name": name, "type": 'PTR'}})
+            except socket.herror:
+                pass
+
     def run(self):
-        scan_result = self.scan()
-        scan_result.pop("stats")
-        scan_result.pop("runtime")
-        for k,v in scan_result.items():
+        self.scan()
+        self.dns_resolution()
+        for k,v in self.scan().items():
             try:
                 self.hosts.append((
                     k,
@@ -30,4 +46,3 @@ class Nmap(object):
                     k,
                     self.unknown
                 ))
-