1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

use crate::errors::*;
use crate::db::{Database, Table};
use crate::schema::*;
use crate::models::*;
use chrono::{NaiveDateTime, Duration, Utc};
use diesel;
use diesel::prelude::*;


#[derive(Identifiable, Queryable, AsChangeset, PartialEq, Debug)]
#[table_name="ttls"]
pub struct Ttl {
    pub id: i32,
    pub family: String,
    pub key: i32,
    pub expire: NaiveDateTime,
}

#[derive(Insertable)]
#[table_name="ttls"]
pub struct NewTtl<'a> {
    pub family: &'a str,
    pub key: i32,
    pub expire: NaiveDateTime,
}

impl Ttl {
    pub fn new(obj: &Insert, key: i32, expire: NaiveDateTime) -> NewTtl {
        NewTtl {
            family: obj.table(),
            key,
            expire,
        }
    }

    pub fn find(obj: &Insert, my_key: i32, db: &Database) -> Result<Option<Ttl>> {
        use crate::schema::ttls::dsl::*;

        ttls
            .filter(family.eq(obj.table()))
            .filter(key.eq(my_key))
            .first::<Self>(db.db())
            .optional()
            .map_err(Error::from)
    }

    pub fn expired(db: &Database) -> Result<Vec<Ttl>> {
        use crate::schema::ttls::dsl::*;

        ttls
            .filter(expire.lt(Self::ttl_to_datetime(0)))
            .load::<Self>(db.db())
            .map_err(Error::from)
    }

    fn ttl_to_datetime(ttl: i32) -> NaiveDateTime {
        // TODO: maybe create Duration from string
        let expire_at = Utc::now() + Duration::seconds(ttl as i64);
        expire_at.naive_utc()
    }

    pub fn create(obj: &Insert, key: i32, ttl: i32, db: &Database) -> Result<()> {
        debug!("Creating ttl on record");
        let expire = Self::ttl_to_datetime(ttl);

        diesel::insert_into(ttls::table)
            .values(NewTtl {
                family: obj.table(),
                key,
                expire,
            })
            .execute(db.db())?;

        Ok(())
    }

    pub fn bump(obj: &Insert, my_key: i32, ttl: i32, db: &Database) -> Result<()> {
        use crate::schema::ttls::dsl::*;

        debug!("Updating ttl on record");

        if let Some(mut old) = Self::find(obj, my_key, db)? {
            let new_expire = Self::ttl_to_datetime(ttl);

            if old.expire < new_expire {
                debug!("Bumping old expire date");

                old.expire = new_expire;

                diesel::update(ttls.filter(id.eq(old.id)))
                    .set(old)
                    .execute(db.db())?;
            }
        } else {
            debug!("Existing record doesn't expire, not setting a ttl");
        }

        Ok(())
    }

    pub fn delete(&self, db: &Database) -> Result<()> {
        let family = self.family.parse::<Table>()?;
        match family {
            Table::Domains => Domain::delete_id(db, self.key)?,
            Table::Subdomains => Subdomain::delete_id(db, self.key)?,
            Table::Ipaddrs => IpAddr::delete_id(db, self.key)?,
            Table::SubdomainIpaddrs => SubdomainIpAddr::delete_id(db, self.key)?,
            Table::Urls => Url::delete_id(db, self.key)?,
            Table::Emails => Email::delete_id(db, self.key)?,
            Table::Phonenumbers => PhoneNumber::delete_id(db, self.key)?,
            Table::Devices => Device::delete_id(db, self.key)?,
            Table::Networks => Network::delete_id(db, self.key)?,
            Table::NetworkDevices => NetworkDevice::delete_id(db, self.key)?,
            Table::Accounts => Account::delete_id(db, self.key)?,
            Table::Breaches => Breach::delete_id(db, self.key)?,
            Table::BreachEmails => BreachEmail::delete_id(db, self.key)?,
            Table::Images => Image::delete_id(db, self.key)?,
            Table::Ports => Port::delete_id(db, self.key)?,
            Table::Netblocks => Netblock::delete_id(db, self.key)?,
            Table::Cryptoaddrs => CryptoAddr::delete_id(db, self.key)?,
        };

        diesel::delete(self)
                .execute(db.db())?;

        Ok(())
    }
}

pub fn reap_expired(db: &Database) -> Result<()> {
    debug!("Reaping expired entities");

    for expired in Ttl::expired(db)? {
        debug!("Expired: {:?}", expired);
        expired.delete(db)?;
    }

    debug!("Finished reaping expired entities");
    Ok(())
}