1 files changed, 26 insertions, 11 deletions
diff --git a/src/db/sql.py b/src/db/sql.py
index f2a9ee2..fc20e36 100644
--- a/src/db/sql.py
+++ b/src/db/sql.py
@@ -1,9 +1,9 @@
 import datetime
-import os
 import sys
 from contextlib import contextmanager
+import os
 
-from sqlalchemy import (Boolean, Column, Date, Integer, String, Text,
+from sqlalchemy import (Boolean, Column, Date, Integer, Serial, String, Text,
                         create_engine, text)
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker
@@ -15,8 +15,9 @@ metadata = Base.metadata
 class Log(Base):
     __tablename__ = "log"
 
-    id = Column(Integer, primary_key=True)
-    timestamp = Column(Date, nullable=False, default=datetime.datetime.utcnow)
+    id = Column(Serial, primary_key=True)
+    timestamp = Column(Date, nullable=False,
+                       default=datetime.datetime.utcnow)
     username = Column(Text, nullable=False)
     logtext = Column(Text, nullable=False)
 
@@ -42,17 +43,20 @@ class Log(Base):
 
 
 class Scanner(Base):
-    __tablename__ = "scanner"
+    __tablename__ = 'scanner'
 
-    id = Column(Integer, primary_key=True)
-    runner = Column(Text, nullable=False, default="*")
-    name = Column(String(128), nullable=False, unique=True)
+    id = Column(Integer, primary_key=True, server_default=text(
+        "nextval('scanner_id_seq'::regclass)"))
+    runner = Column(Text, server_default=text("'*'::text"))
+    name = Column(String(128), nullable=False)
     active = Column(Boolean, nullable=False)
-    interval = Column(Integer, nullable=False, server_default=text("300"))
+    interval = Column(Integer, nullable=False,
+                      server_default=text("300"))
     starttime = Column(Date)
-    hostname = Column(String(128), nullable=False, unique=True)
+    endtime = Column(Date)
+    maxruns = Column(Integer, server_default=text("1"))
+    hostname = Column(String(128), nullable=False)
     port = Column(Integer, nullable=False)
-    maxruns = Column(Integer, nullable=False, default=1)
 
     def as_dict(self):
         d = {}
@@ -107,6 +111,17 @@ class Scanner(Base):
                     results.append(scanner.as_dict())
         return results
 
+    @classmethod
+    def edit(cls, name, active):
+        with sqla_session() as session:
+            scanners = session.query(Scanner).filter(
+                Scanner.name == name).all()
+            if not scanners:
+                return None
+            for scanner in scanners:
+                scanner.active = active
+        return True
+
 
 def get_sqlalchemy_conn_str(**kwargs) -> str:
     try: