Refactor code to use SQLAlchemy instead of PeeWee.

db.py

@@ -1,33 +1,40 @@
 from datetime import datetime
-import peewee
+from sqlalchemy import create_engine, String, Integer, Column, DateTime
+from sqlalchemy.orm import declarative_base, sessionmaker
+import sqlalchemy.exc
 
+# Create a new database engine, in this case, a sqlite database
+engine = create_engine("sqlite:///devices.db")
+# Create a new table with a name, count, amount, and valid column
+Base = declarative_base()
 
-db = peewee.SqliteDatabase("devices.db")
 
-
-class Device(peewee.Model):
+class Devices(Base):
     """
-    Database model for devices.
-
-    Args:
-        peewee.Model: The base model class that other models will inherit from.
+    The Devices class is a table that stores information about devices.
 
-    Returns:
-        None
+    Attributes:
+        id: The id of the device.
+        ip: The IP address of the device.
+        hostname: The hostname of the device.
+        device_type: The device type.
+        date_updated: The date the device was updated.
     """
 
-    id = peewee.IntegerField(primary_key=True)
-    ip = peewee.IPField(unique=True)
-    hostname = peewee.TextField()
-    date_added = peewee.DateTimeField(default=datetime.now)
-    device_type = peewee.CharField()
+    __tablename__ = "devices"
+    id = Column(Integer, primary_key=True)
+    ip = Column(String, unique=True, nullable=False)
+    hostname = Column(String, nullable=False)
+    device_type = Column(String, nullable=False)
+    date_updated = Column(DateTime, default=datetime.now(), onupdate=datetime.now())
 
-    class Meta:
-        database = db  # This model uses the "devices.db" database.
 
-
-db.connect()
-db.create_tables([Device])
+# Create the table
+Base.metadata.create_all(engine)
+# Create a session to use the tables, bound to above engine
+Session = sessionmaker(bind=engine)
+# Create a session
+session = Session()
 
 
 def insert_device(ip, hostname, device_type):
@@ -40,68 +47,89 @@ def insert_device(ip, hostname, device_type):
         device_type: The device type.
 
     Returns:
-        None
+        True if the device was added, False if the device already exists.
     """
+
     try:
-        device = Device(ip=ip, hostname=hostname, device_type=device_type)
-        device.save()
-    except peewee.IntegrityError:
-        return 0
+        device = Devices(ip=ip, hostname=hostname, device_type=device_type)
+        session.add(device)
+        session.commit()
+    except sqlalchemy.exc.IntegrityError:
+        return False
+    else:
+        return True
 
 
-def list_all_ips():
+def insert_device_bulk(devices_object):
     """
-    List all the IPs in the database.
+    Insert multiple device into the database.
 
     Args:
-        None
+        devices_object: an object containing the devices to be added.
 
     Returns:
-        A list of all the IPs in the database with their device IDs.
+        None
     """
 
-    device = []
-    device_id = []
-    for device_ip in Device.select():
-        device.append(device_ip.ip)
-        device_id.append(device_ip.id)
-    return device, device_id
+    final_list = []
+    print("Checking for duplicate devices in the database ...")
+    for xx in devices_object:
+        xx = Devices(ip=xx.ip, hostname=xx.hostname, device_type=xx.device_type)
+        final_list.append(xx)
+
+    print("Adding devices to the database.")
+    session.add_all(final_list)
+    # Commit the changes
+    session.commit()
 
 
 def list_all_ips_with_type():
     """
     List all the IPs in the database with their device types.
 
-    Args:
-        None
-
-
     Returns:
         A list of all the IPs in the database with their device types.
     """
 
-    device_id = []
-    device_ip = []
-    device_hostname = []
-    device_type = []
+    devices = session.query(Devices).all()
+    return devices
+
+
+def remove_device(device_id):
+    """
+    Remove a device from the database.
+
+    Args:
+        device_id: The id of the device to remove.
 
-    for device in Device.select():
-        device_id.append(device.id)
-        device_ip.append(device.ip)
-        device_hostname.append(device.hostname)
-        device_type.append(device.device_type)
-    return device_id, device_ip, device_hostname, device_type
+    Returns:
+        True if the device was removed, False if the device was not found.
+    """
 
+    try:
+        device = session.query(Devices).filter(Devices.id == device_id).one()
+        session.delete(device)
+        session.commit()
+    except sqlalchemy.exc.NoResultFound:
+        return False
+    else:
+        return True
 
-def delete_device(id):
+
+def is_device_in_db(ip):
     """
-    Delete a device from the database.
+    Check if a device is in the database.
 
     Args:
-        id: The ID of the device.
+        ip: The IP of the device to check.
 
     Returns:
-        1 if the device was deleted, 0 if the device was not deleted.
+        True if the device is in the database, False if the device is not in the database.
     """
-    device = Device.delete_by_id(id)
-    return device
+
+    device = session.query(Devices).filter(Devices.ip == ip).one_or_none()
+
+    if device is None:
+        return False
+    else:
+        return True

device.py

@@ -19,8 +19,8 @@ def detect_device(ip):
         "username": my_secrets.USERNAME,
         "password": my_secrets.PASSWORD,
         "global_delay_factor": 4,
-        "banner_timeout": 200,
-        "conn_timeout": 200,
+        "banner_timeout": 1000,
+        "conn_timeout": 1000,
     }
 
     guesser = SSHDetect(**device)

manage.py

@@ -1,15 +1,27 @@
 # manage.py
 
 import argparse
-import sys
 import db
 import get_config
 import utils
 import conf
 import hosts
 import device
 import get_hostname
+import time
 
+start_time = time.perf_counter()
+
+
+# Class to hold device information to be sent to the db
+class Item:
+    def __init__(self, d_ip, d_hostname, d_type):
+        self.ip = d_ip
+        self.hostname = d_hostname
+        self.device_type = d_type
+
+
+# initialize the parser
 parser = argparse.ArgumentParser(
     usage="test.py [options]", description="Switch management tool."
 )
@@ -53,57 +65,74 @@
 if args.scan:
     alive_hosts = hosts.scan_cidr(args.scan)
     alive_hosts_filtered = []
+    alive_hosts_final = []
     if args.add is not None:
+        # If -skip flag is used, remove the IPs from the list.
         if args.skip is not None:
             skip_hosts = args.skip.split(",")
             skip_hosts = [value.strip() for value in skip_hosts]
             alive_hosts_filtered = [x for x in alive_hosts if x not in set(skip_hosts)]
         else:
             alive_hosts_filtered = alive_hosts
 
+        # check if device already exists in the database
         for device_ip in alive_hosts_filtered:
-            device_type = device.detect_device(device_ip)
-            hostname = get_hostname.get_hostname(
-                device_ip=device_ip, device_type=device_type
-            )
-            print(f"Processing {device_ip} ...")
-            if (
-                db.insert_device(
-                    ip=device_ip, hostname=hostname, device_type=device_type
-                )
-                == 0
-            ):
-                print(f"{device_ip} already exists in the database.")
+            if db.is_device_in_db(device_ip) is False:
+                alive_hosts_final.append(device_ip)
             else:
-                print(f"{device_ip} has been added to the database.")
+                print(f"Device with ip {device_ip} already exists in the database.")
+
+        devices_object = []
+
+        # get device type and hostname for each device, then add to the database
+        if len(alive_hosts_final) != 0:
+            for device_ip in alive_hosts_final:
+                print(f"Detecting {device_ip} device type ...")
+                device_type = device.detect_device(device_ip)
+                print(f"Getting hostname for {device_ip} ...")
+                hostname = get_hostname.get_hostname(
+                    device_ip=device_ip, device_type=device_type
+                )
+                single_device = Item(device_ip, hostname, device_type)
+                devices_object.append(single_device)
+            db.insert_device_bulk(devices_object)
+        else:
+            print("No devices to add to the database.")
+
 elif args.list:
-    if (len(db.list_all_ips()[0])) == 0:
-        print("No devices in database.")
+    devices = db.list_all_ips_with_type()
+    if len(devices) == 0:
+        print("No devices in the database.")
     else:
-        for device_id, device_ip, device_hostname, device_type in zip(
-            db.list_all_ips_with_type()[0],
-            db.list_all_ips_with_type()[1],
-            db.list_all_ips_with_type()[2],
-            db.list_all_ips_with_type()[3],
-        ):
-            print(f"{device_id} : {device_ip} : {device_hostname} : {device_type}")
+        for device in devices:
+            print(
+                f"{device.id} | {device.ip} | {device.hostname} | {device.device_type}"
+            )
+
 elif args.add:
-    hostname = device.get_hostname(device_ip=args.add, device_type=args.type)
+    print(f"Detecting {args.add} device type ...")
+    device_type = device.detect_device(args.add)
+    print(f"Getting hostname for {args.add} ...")
+    hostname = get_hostname.get_hostname(device_ip=args.add, device_type=device_type)
 
-    if args.type is None:
-        print("You must specify a device type with the -type flag.")
+    if db.insert_device(args.add, hostname, device_type) is True:
+        print(f"Device with ip {args.add} has been added to the database.")
     else:
-        if db.insert_device(ip=args.add, hostname=hostname, device_type=args.type) == 0:
-            print(f"{args.add} already exists in the database.")
-        else:
-            print(f"{args.add} has been added to the database.")
+        print(f"Device with ip {args.add} already exists in the database.")
+
 elif args.remove:
-    if db.delete_device(id=args.remove) == 1:
-        print("Device removed.")
+    if db.remove_device(args.remove) is True:
+        print(f"Device with id {args.remove} has been removed from the database.")
     else:
-        print("Device ID does not exist.")
+        print(f"Device with id {args.remove} does not exist in the database.")
+
 elif args.fetchall:
     get_config.fetch_all_configs()
+
 elif args.clean:
     print("Cleaning up running-configs/ directory.")
     utils.del_oldest_configs(conf.MAX_CONFIGS)
+
+end_time = time.perf_counter()
+elapsed_time = end_time - start_time
+print("Execution time:", elapsed_time, "seconds")

requirements.txt

@@ -1,5 +1,5 @@
-peewee~=3.17.1
 netmiko~=4.3.0
 tqdm~=4.66.2
 icmplib~=3.0.4
-argparse~=1.4.0
+argparse~=1.4.0
+SQLAlchemy~=2.0.27