# database.py is part of Panopticon.

# Panopticon is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# Panopticon is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with Panopticon.  If not, see <http://www.gnu.org/licenses/>.

from contextlib import contextmanager
from paramiko import RSAKey as pRSAKey, DSSKey
from sqlalchemy import create_engine, Column, DateTime, String, Integer, Text, Boolean
from sqlalchemy.orm import sessionmaker, relationship, backref
from sqlalchemy.sql import not_ 
from sqlalchemy.schema import ForeignKey
from sqlalchemy.pool import NullPool
from sqlalchemy.ext.declarative import declarative_base
from panopticon.core.util.database import key_value_property

Base = declarative_base()

class Value(Base):
    __tablename__ = "values"
    id = Column(Integer(), primary_key=True)
    name = Column(String(1000))
    value = Column(String(1000), nullable=True)
    parent_id = Column(Integer, ForeignKey("values.id"), nullable=True)
    values = relationship("Value", backref=backref('parent', remote_side=[id], 
        cascade="all"))
    type = Column(String(20))

    def __init__(self, name, _type, value="", parent_id=None):
        self.name = name
        self.type = _type
        self.value = value
        self.parent_id = parent_id

    @property
    def root(self):
        return self.id == self.parent

class Service(Base):
    __tablename__ = "services"
    name = Column(String(50), primary_key=True)

class Computer(Base):
    __tablename__ = "computers"
    __table_args__ = {'sqlite_autoincrement':True}

    name = Column(String(255), primary_key=True)
    key_name = Column(String(100), ForeignKey('keys.name', onupdate="CASCADE"))
    active = Column(Boolean(), default=True)
    key = relationship("Key", backref=backref('computers'))
    logs = relationship("Log", backref="computer", order_by="Log.time")

    def __init__(self, name, key_name="", active=True):
        self.name = name
        self.active = active
        self.key_name = key_name

class Log(Base):
    __tablename__ = "logs"
    id = Column('id', Integer, primary_key=True)
    time = Column(DateTime())
    level = Column(String(10))
    message = Column(Text())
    computer_name = Column(String(255), ForeignKey('computers.name',
        ondelete="CASCADE", onupdate="CASCADE"), index=True)
    service_name = Column(String(255), ForeignKey('services.name',
        ondelete="CASCADE", onupdate="CASCADE"), index=True)
    role_name = Column(String(255), index=True)
    action_name = Column(String(255), index=True)

    def __init__(self, time, level, message, computer_name="", 
            service_name="", role_name="", action_name=""):
        self.time = time
        self.level = level
        self.message = message
        self.computer_name = computer_name

class FileTrack(Base):
    __tablename__ = "filetracks"
    uid = Column("uid", String(32), primary_key=True)
    _computer_name = Column("computer_name", String(255),ForeignKey('computers.name'))
    _path = Column("path", Text())
    modification_time = Column("modification_time", DateTime())
    md5 = Column("md5", String(32))

    def __init__(self, computer_name, path, modification_time, md5=""):
        self.computer_name = computer_name
        self.path = path
        self.modification_time = modification_time
        self.md5 = md5
        self.update_uid()

    @property
    def computer_name(self):
        return self._computer_name

    @computer_name.setter
    def computer_name(self, value):
        self._computer_name = value
        self.update_uid()

    @property
    def path(self):
        return self._path

    @path.setter
    def path(self, value):
        self._path = value
        self.update_uid()

    def update_uid(self):
        if self.computer_name and self.path:
            self.uid = "%s:%s" % (self.computer_name, self.path)
        else:
            self.uid = ""

class Key(Base):
    __tablename__ = "keys"
    name = Column(String(100), primary_key=True)
    algorithm = Column(String(20))
    v1 = Column(String(2048))
    v2 = Column(String(2048))
    v3 = Column(String(2048))
    v4 = Column(String(2048))
    key_class = None
    key_vals = []
    __mapper_args__ = {'polymorphic_on' : algorithm}

    @classmethod
    def build_from_paramiko_key(cls, name, p_key):
        if isinstance(p_key, pRSAKey):
            return RSAKey(name, p_key.e, p_key.n)
        elif isinstance(p_key, DSSKey):
            return DSAKey(name, p_key.p, p_key.q, p_key.g, p_key.y)
        else:
            raise Exception("Not valid key")

    def __init__(self, name, algorithm, v1, v2, v3, v4):
        self.name = name
        self.algorithm = algorithm
        self.v1 = v1
        self.v2 = v2
        self.v3 = v3
        self.v4 = v4

    def get_paramiko_key(self):
        vals = [ getattr(self, x) for x in self.key_vals ]
        return self.key_class(vals=vals)

class RSAKey(Key):
    __mapper_args__ = {'polymorphic_identity':'rsa'}
    key_class = pRSAKey
    key_vals = [ 'e', 'n' ]

    def __init__(self, name, e, n):
        self.name = name
        self.algorithm = "rsa"
        self.e = e
        self.n = n

    e = key_value_property("v1")
    n = key_value_property("v2")


class DSAKey(Key):
    __mapper_args__ = {'polymorphic_identity':'dsa'}
    key_class = DSSKey
    key_vals = [ 'p', 'q', 'g', 'y' ]

    def __init__(self, name, p, q, g, y):
        self.name = name
        self.algorithm = "dsa"
        self.p = p
        self.q = q
        self.g = g
        self.y = y

    p = key_value_property("v1")
    q = key_value_property("v2")
    g = key_value_property("v3")
    y = key_value_property("v4")


class PanopticonDB(object):
    def __init__(self, panopticon, engine=None):
        self.panopticon = panopticon
        self.engine = engine if engine is not None else create_engine(panopticon.db_url, poolclass=NullPool)
        Base.metadata.create_all(self.engine)
        self.Session = sessionmaker(bind=self.engine)
        self.sync()

    @contextmanager
    def get_session(self):
        session = self.Session()
        yield session
        session.commit()
        session.close()

    def purge(self,sure=False):
        if sure:
            Base.metadata.drop_all(self.engine)
            Base.metadata.create_all(self.engine)

    def sync(self):
        computer_names = [ x[0] for x in self.panopticon.computers ]
        with self.get_session() as session:
            session.execute(Computer.__table__.update().where(Computer.name.in_(computer_names)).values(active=True))
            session.execute(Computer.__table__.update().where(not_(Computer.name.in_(computer_names))).values(active=True))
        
