From 42ca33f913c13bf5e905ba480ca3e87a31968176 Mon Sep 17 00:00:00 2001 From: kpcyrd Date: Wed, 12 Sep 2018 12:54:17 +0200 Subject: [PATCH] Add database code --- Cargo.lock | 137 ++++++++++++++++++ Cargo.toml | 2 + diesel.toml | 5 + migrations/.gitkeep | 0 .../2018-09-12-102647_create_domains/down.sql | 5 + .../2018-09-12-102647_create_domains/up.sql | 40 +++++ src/errors.rs | 2 + src/main.rs | 14 +- src/migrations.rs | 11 ++ src/models.rs | 43 ++++++ src/schema.rs | 53 +++++++ src/shell.rs | 28 +++- src/worker.rs | 41 +++++- 13 files changed, 376 insertions(+), 5 deletions(-) create mode 100644 diesel.toml create mode 100644 migrations/.gitkeep create mode 100644 migrations/2018-09-12-102647_create_domains/down.sql create mode 100644 migrations/2018-09-12-102647_create_domains/up.sql create mode 100644 src/errors.rs create mode 100644 src/migrations.rs create mode 100644 src/models.rs create mode 100644 src/schema.rs diff --git a/Cargo.lock b/Cargo.lock index dc17a494..5338637b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -50,6 +50,11 @@ name = "bitflags" version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "byteorder" +version = "1.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "cc" version = "1.0.24" @@ -90,6 +95,35 @@ dependencies = [ "lazy_static 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "diesel" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "byteorder 1.2.6 (registry+https://github.com/rust-lang/crates.io-index)", + "diesel_derives 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)", + "libsqlite3-sys 0.9.3 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "diesel_derives" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "proc-macro2 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)", + "quote 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)", + "syn 0.13.11 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "diesel_migrations" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "migrations_internals 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)", + "migrations_macros 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "dirs" version = "1.0.3" @@ -203,6 +237,15 @@ name = "libc" version = "0.2.43" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "libsqlite3-sys" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "pkg-config 0.3.14 (registry+https://github.com/rust-lang/crates.io-index)", + "vcpkg 0.2.6 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "log" version = "0.4.5" @@ -234,6 +277,24 @@ dependencies = [ "libc 0.2.43 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "migrations_internals" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "diesel 1.3.2 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "migrations_macros" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "migrations_internals 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)", + "quote 0.3.15 (registry+https://github.com/rust-lang/crates.io-index)", + "syn 0.11.11 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "nix" version = "0.11.0" @@ -256,6 +317,14 @@ name = "pkg-config" version = "0.3.14" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "proc-macro2" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "unicode-xid 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "proc-macro2" version = "0.4.19" @@ -281,6 +350,19 @@ name = "quick-error" version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "quote" +version = "0.3.15" +source = "registry+https://github.com/rust-lang/crates.io-index" + +[[package]] +name = "quote" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "proc-macro2 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "quote" version = "0.6.8" @@ -394,6 +476,8 @@ name = "sn0int" version = "0.1.0" dependencies = [ "colored 1.6.1 (registry+https://github.com/rust-lang/crates.io-index)", + "diesel 1.3.2 (registry+https://github.com/rust-lang/crates.io-index)", + "diesel_migrations 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)", "env_logger 0.5.13 (registry+https://github.com/rust-lang/crates.io-index)", "failure 0.1.2 (registry+https://github.com/rust-lang/crates.io-index)", "hlua-badtouch 0.4.2 (registry+https://github.com/rust-lang/crates.io-index)", @@ -430,6 +514,26 @@ dependencies = [ "syn 0.14.9 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "syn" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "quote 0.3.15 (registry+https://github.com/rust-lang/crates.io-index)", + "synom 0.11.3 (registry+https://github.com/rust-lang/crates.io-index)", + "unicode-xid 0.0.4 (registry+https://github.com/rust-lang/crates.io-index)", +] + +[[package]] +name = "syn" +version = "0.13.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "proc-macro2 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)", + "quote 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)", + "unicode-xid 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "syn" version = "0.14.9" @@ -440,6 +544,14 @@ dependencies = [ "unicode-xid 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)", ] +[[package]] +name = "synom" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +dependencies = [ + "unicode-xid 0.0.4 (registry+https://github.com/rust-lang/crates.io-index)", +] + [[package]] name = "synstructure" version = "0.9.0" @@ -513,6 +625,11 @@ name = "unicode-width" version = "0.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "unicode-xid" +version = "0.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "unicode-xid" version = "0.1.0" @@ -538,6 +655,11 @@ name = "utf8parse" version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" +[[package]] +name = "vcpkg" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" + [[package]] name = "vec_map" version = "0.8.1" @@ -596,11 +718,15 @@ dependencies = [ "checksum backtrace 0.3.9 (registry+https://github.com/rust-lang/crates.io-index)" = "89a47830402e9981c5c41223151efcced65a0510c13097c769cede7efb34782a" "checksum backtrace-sys 0.1.24 (registry+https://github.com/rust-lang/crates.io-index)" = "c66d56ac8dabd07f6aacdaf633f4b8262f5b3601a810a0dcddffd5c22c69daa0" "checksum bitflags 1.0.4 (registry+https://github.com/rust-lang/crates.io-index)" = "228047a76f468627ca71776ecdebd732a3423081fcf5125585bcd7c49886ce12" +"checksum byteorder 1.2.6 (registry+https://github.com/rust-lang/crates.io-index)" = "90492c5858dd7d2e78691cfb89f90d273a2800fc11d98f60786e5d87e2f83781" "checksum cc 1.0.24 (registry+https://github.com/rust-lang/crates.io-index)" = "70f2a88c2e69ceee91c209d8ef25b81fc1a65f42c7f14dfd59d1fed189e514d1" "checksum cfg-if 0.1.5 (registry+https://github.com/rust-lang/crates.io-index)" = "0c4e7bb64a8ebb0d856483e1e682ea3422f883c5f5615a90d51a2c82fe87fdd3" "checksum clap 2.32.0 (registry+https://github.com/rust-lang/crates.io-index)" = "b957d88f4b6a63b9d70d5f454ac8011819c6efa7727858f458ab71c756ce2d3e" "checksum cloudabi 0.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "ddfc5b9aa5d4507acaf872de71051dfd0e309860e88966e1051e462a077aac4f" "checksum colored 1.6.1 (registry+https://github.com/rust-lang/crates.io-index)" = "dc0a60679001b62fb628c4da80e574b9645ab4646056d7c9018885efffe45533" +"checksum diesel 1.3.2 (registry+https://github.com/rust-lang/crates.io-index)" = "e71e7a348ae6064e86c4cf0709f0e4c3ef6f30e8e7d3dc05737164af4ebd3511" +"checksum diesel_derives 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "03bcaf77491f53e400d5ee3bdd57142ea4e1c47fe9217b3361ff9a76ca0e3d37" +"checksum diesel_migrations 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "17b42c35d1ce9e8d57a3e7001b4127f2bc1b073a89708bb7019f5be27c991c28" "checksum dirs 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "f679c09c1cf5428702cc10f6846c56e4e23420d3a88bcc9335b17c630a7b710b" "checksum env_logger 0.5.13 (registry+https://github.com/rust-lang/crates.io-index)" = "15b0a4d2e39f8420210be8b27eeda28029729e2fd4291019455016c348240c38" "checksum error-chain 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)" = "ff511d5dc435d703f4971bc399647c9bc38e20cb41452e3b9feb4765419ed3f3" @@ -615,16 +741,22 @@ dependencies = [ "checksum lazy_static 0.2.11 (registry+https://github.com/rust-lang/crates.io-index)" = "76f033c7ad61445c5b347c7382dd1237847eb1bce590fe50365dcb33d546be73" "checksum lazy_static 1.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "ca488b89a5657b0a2ecd45b95609b3e848cf1755da332a0da46e2b2b1cb371a7" "checksum libc 0.2.43 (registry+https://github.com/rust-lang/crates.io-index)" = "76e3a3ef172f1a0b9a9ff0dd1491ae5e6c948b94479a3021819ba7d860c8645d" +"checksum libsqlite3-sys 0.9.3 (registry+https://github.com/rust-lang/crates.io-index)" = "d3711dfd91a1081d2458ad2d06ea30a8755256e74038be2ad927d94e1c955ca8" "checksum log 0.4.5 (registry+https://github.com/rust-lang/crates.io-index)" = "d4fcce5fa49cc693c312001daf1d13411c4a5283796bac1084299ea3e567113f" "checksum lua52-sys 0.1.1 (registry+https://github.com/rust-lang/crates.io-index)" = "39c2835cc62c821132f91bb28fa41ae178f56b056ab9577c6f7ed81a908b1187" "checksum matches 0.1.8 (registry+https://github.com/rust-lang/crates.io-index)" = "7ffc5c5338469d4d3ea17d269fa8ea3512ad247247c30bd2df69e68309ed0a08" "checksum memchr 2.0.2 (registry+https://github.com/rust-lang/crates.io-index)" = "a3b4142ab8738a78c51896f704f83c11df047ff1bda9a92a661aa6361552d93d" +"checksum migrations_internals 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "8cf7c8c4f83fa9f47440c0b4af99973502de55e6e7b875f693bd263e03f93e7e" +"checksum migrations_macros 1.3.0 (registry+https://github.com/rust-lang/crates.io-index)" = "79f12499ef7353bdeca2d081bc61edd8351dac09a33af845952009b5a3d68c1a" "checksum nix 0.11.0 (registry+https://github.com/rust-lang/crates.io-index)" = "d37e713a259ff641624b6cb20e3b12b2952313ba36b6823c0f16e6cfd9e5de17" "checksum percent-encoding 1.0.1 (registry+https://github.com/rust-lang/crates.io-index)" = "31010dd2e1ac33d5b46a5b413495239882813e0369f8ed8a5e266f173602f831" "checksum pkg-config 0.3.14 (registry+https://github.com/rust-lang/crates.io-index)" = "676e8eb2b1b4c9043511a9b7bea0915320d7e502b0a079fb03f9635a5252b18c" +"checksum proc-macro2 0.3.8 (registry+https://github.com/rust-lang/crates.io-index)" = "1b06e2f335f48d24442b35a19df506a835fb3547bc3c06ef27340da9acf5cae7" "checksum proc-macro2 0.4.19 (registry+https://github.com/rust-lang/crates.io-index)" = "ffe022fb8c8bd254524b0b3305906c1921fa37a84a644e29079a9e62200c3901" "checksum publicsuffix 1.5.1 (registry+https://github.com/rust-lang/crates.io-index)" = "e0eafc29e6abf231d6a2999608b35d4397f212a77240ab5ac5f65364f5bea6c2" "checksum quick-error 1.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "9274b940887ce9addde99c4eee6b5c44cc494b182b97e73dc8ffdcb3397fd3f0" +"checksum quote 0.3.15 (registry+https://github.com/rust-lang/crates.io-index)" = "7a6e920b65c65f10b2ae65c831a81a073a89edd28c7cce89475bff467ab4167a" +"checksum quote 0.5.2 (registry+https://github.com/rust-lang/crates.io-index)" = "9949cfe66888ffe1d53e6ec9d9f3b70714083854be20fd5e271b232a017401e8" "checksum quote 0.6.8 (registry+https://github.com/rust-lang/crates.io-index)" = "dd636425967c33af890042c483632d33fa7a18f19ad1d7ea72e8998c6ef8dea5" "checksum rand 0.5.5 (registry+https://github.com/rust-lang/crates.io-index)" = "e464cd887e869cddcae8792a4ee31d23c7edd516700695608f5b98c67ee0131c" "checksum rand_core 0.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "edecf0f94da5551fc9b492093e30b041a891657db7940ee221f9d2f66e82eef2" @@ -640,7 +772,10 @@ dependencies = [ "checksum strsim 0.7.0 (registry+https://github.com/rust-lang/crates.io-index)" = "bb4f380125926a99e52bc279241539c018323fab05ad6368b56f93d9369ff550" "checksum structopt 0.2.10 (registry+https://github.com/rust-lang/crates.io-index)" = "d8e9ad6a11096cbecdcca0cc6aa403fdfdbaeda2fb3323a39c98e6a166a1e45a" "checksum structopt-derive 0.2.10 (registry+https://github.com/rust-lang/crates.io-index)" = "4cbce8ccdc62166bd594c14396a3242bf94c337a51dbfa9be1076dd74b3db2af" +"checksum syn 0.11.11 (registry+https://github.com/rust-lang/crates.io-index)" = "d3b891b9015c88c576343b9b3e41c2c11a51c219ef067b264bd9c8aa9b441dad" +"checksum syn 0.13.11 (registry+https://github.com/rust-lang/crates.io-index)" = "14f9bf6292f3a61d2c716723fdb789a41bbe104168e6f496dc6497e531ea1b9b" "checksum syn 0.14.9 (registry+https://github.com/rust-lang/crates.io-index)" = "261ae9ecaa397c42b960649561949d69311f08eeaea86a65696e6e46517cf741" +"checksum synom 0.11.3 (registry+https://github.com/rust-lang/crates.io-index)" = "a393066ed9010ebaed60b9eafa373d4b1baac186dd7e008555b0f702b51945b6" "checksum synstructure 0.9.0 (registry+https://github.com/rust-lang/crates.io-index)" = "85bb9b7550d063ea184027c9b8c20ac167cd36d3e06b3a40bceb9d746dc1a7b7" "checksum termcolor 1.0.3 (registry+https://github.com/rust-lang/crates.io-index)" = "ff3bac0e465b59f194e7037ed404b0326e56ff234d767edc4c5cc9cd49e7a2c7" "checksum termion 1.5.1 (registry+https://github.com/rust-lang/crates.io-index)" = "689a3bdfaab439fd92bc87df5c4c78417d3cbe537487274e9b0b2dce76e92096" @@ -651,10 +786,12 @@ dependencies = [ "checksum unicode-normalization 0.1.7 (registry+https://github.com/rust-lang/crates.io-index)" = "6a0180bc61fc5a987082bfa111f4cc95c4caff7f9799f3e46df09163a937aa25" "checksum unicode-segmentation 1.2.1 (registry+https://github.com/rust-lang/crates.io-index)" = "aa6024fc12ddfd1c6dbc14a80fa2324d4568849869b779f6bd37e5e4c03344d1" "checksum unicode-width 0.1.5 (registry+https://github.com/rust-lang/crates.io-index)" = "882386231c45df4700b275c7ff55b6f3698780a650026380e72dabe76fa46526" +"checksum unicode-xid 0.0.4 (registry+https://github.com/rust-lang/crates.io-index)" = "8c1f860d7d29cf02cb2f3f359fd35991af3d30bac52c57d265a3c461074cb4dc" "checksum unicode-xid 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "fc72304796d0818e357ead4e000d19c9c174ab23dc11093ac919054d20a6a7fc" "checksum url 1.7.1 (registry+https://github.com/rust-lang/crates.io-index)" = "2a321979c09843d272956e73700d12c4e7d3d92b2ee112b31548aef0d4efc5a6" "checksum utf8-ranges 1.0.1 (registry+https://github.com/rust-lang/crates.io-index)" = "fd70f467df6810094968e2fce0ee1bd0e87157aceb026a8c083bcf5e25b9efe4" "checksum utf8parse 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "a15ea87f3194a3a454c78d79082b4f5e85f6956ddb6cb86bbfbe4892aa3c0323" +"checksum vcpkg 0.2.6 (registry+https://github.com/rust-lang/crates.io-index)" = "def296d3eb3b12371b2c7d0e83bfe1403e4db2d7a0bba324a12b21c4ee13143d" "checksum vec_map 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)" = "05c78687fb1a80548ae3250346c3db86a80a7cdd77bda190189f2d0a0987c81a" "checksum version_check 0.1.4 (registry+https://github.com/rust-lang/crates.io-index)" = "7716c242968ee87e5542f8021178248f267f295a5c4803beae8b8b7fd9bc6051" "checksum void 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)" = "6a02e4885ed3bc0f2de90ea6dd45ebcbb66dacffe03547fadbb0eeae2770887d" diff --git a/Cargo.toml b/Cargo.toml index 32cc1267..98bfbc07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,5 +20,7 @@ colored = "1.6" lazy_static = "1.0" shellwords = "0.1" publicsuffix = { version="1.5", default-features=false } +diesel = { version = "1.0.0", features = ["sqlite"] } +diesel_migrations = { version = "1.3.0", features = ["sqlite"] } # https://crates.io/crates/tor_control diff --git a/diesel.toml b/diesel.toml new file mode 100644 index 00000000..92267c82 --- /dev/null +++ b/diesel.toml @@ -0,0 +1,5 @@ +# For documentation on how to configure this file, +# see diesel.rs/guides/configuring-diesel-cli + +[print_schema] +file = "src/schema.rs" diff --git a/migrations/.gitkeep b/migrations/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/migrations/2018-09-12-102647_create_domains/down.sql b/migrations/2018-09-12-102647_create_domains/down.sql new file mode 100644 index 00000000..fc52d242 --- /dev/null +++ b/migrations/2018-09-12-102647_create_domains/down.sql @@ -0,0 +1,5 @@ +DROP TABLE urls; +DROP TABLE subdomain_ipaddrs; +DROP TABLE subdomains; +DROP TABLE ipaddrs; +DROP TABLE domains; diff --git a/migrations/2018-09-12-102647_create_domains/up.sql b/migrations/2018-09-12-102647_create_domains/up.sql new file mode 100644 index 00000000..a86cb995 --- /dev/null +++ b/migrations/2018-09-12-102647_create_domains/up.sql @@ -0,0 +1,40 @@ +CREATE TABLE domains ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + value VARCHAR NOT NULL, + CONSTRAINT domain_unique UNIQUE (value) +); + +CREATE TABLE subdomains ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + domain_id INTEGER NOT NULL, + value VARCHAR NOT NULL, + FOREIGN KEY(domain_id) REFERENCES domains(id), + CONSTRAINT subdomain_unique UNIQUE (value) +); + +/* family maybe not needed */ +CREATE TABLE ipaddrs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + family VARCHAR NOT NULL, + value VARCHAR NOT NULL, + CONSTRAINT ipaddr_unique UNIQUE (value) +); + +CREATE TABLE subdomain_ipaddrs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + subdomain_id INTEGER NOT NULL, + ip_addr_id INTEGER NOT NULL, + FOREIGN KEY(subdomain_id) REFERENCES domains(id), + FOREIGN KEY(ip_addr_id) REFERENCES ipaddrs(id), + CONSTRAINT subdomain_ipaddr_unique UNIQUE (subdomain_id, ip_addr_id) +); + +CREATE TABLE urls ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + subdomain_id INTEGER NOT NULL, + value VARCHAR NOT NULL, + status INTEGER, + body BLOB, + FOREIGN KEY(subdomain_id) REFERENCES domains(id), + CONSTRAINT url_unique UNIQUE (value) +); diff --git a/src/errors.rs b/src/errors.rs new file mode 100644 index 00000000..245cff2e --- /dev/null +++ b/src/errors.rs @@ -0,0 +1,2 @@ +pub use failure::{Error, ResultExt}; +pub type Result = ::std::result::Result; diff --git a/src/main.rs b/src/main.rs index e199c5df..40bd221a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,8 +5,14 @@ extern crate colored; extern crate failure; extern crate shellwords; #[macro_use] extern crate lazy_static; +#[macro_use] extern crate diesel; +#[macro_use] extern crate diesel_migrations; +pub mod errors; pub mod engine; +pub mod migrations; +pub mod models; +pub mod schema; pub mod shell; pub mod term; pub mod worker; @@ -15,5 +21,11 @@ pub mod worker; fn main() { env_logger::init(); - shell::run(); + if let Err(err) = shell::run() { + eprintln!("Error: {}", err); + for cause in err.iter_chain().skip(1) { + eprintln!("Because: {}", cause); + } + std::process::exit(1); + } } diff --git a/src/migrations.rs b/src/migrations.rs new file mode 100644 index 00000000..473d38cf --- /dev/null +++ b/src/migrations.rs @@ -0,0 +1,11 @@ +#![allow(unused_imports)] +use errors::*; + +use diesel::sqlite::*; + +embed_migrations!(); + +pub fn run(conn: &SqliteConnection) -> Result<()> { + embedded_migrations::run(conn)?; + Ok(()) +} diff --git a/src/models.rs b/src/models.rs new file mode 100644 index 00000000..722d20e5 --- /dev/null +++ b/src/models.rs @@ -0,0 +1,43 @@ +use schema::*; + +#[derive(Identifiable, Queryable, PartialEq, Debug)] +#[table_name="domains"] +pub struct Domain { + pub id: i64, + pub value: String, +} + +#[derive(Identifiable, Queryable, Associations, PartialEq, Debug)] +#[belongs_to(Domain)] +#[table_name="subdomains"] +pub struct Subdomain { + pub id: i64, + pub domain_id: i64, + pub value: String, +} + +#[derive(Identifiable, Queryable, Associations, PartialEq, Debug)] +#[table_name="ipaddrs"] +pub struct IpAddr { + pub id: i64, +} + +#[derive(Identifiable, Queryable, Associations)] +#[belongs_to(Subdomain)] +#[belongs_to(IpAddr)] +#[table_name="subdomain_ipaddrs"] +pub struct SubdomainIpAddr { + pub id: i64, + pub subdomain_id: i64, + pub ip_addr_id: i64, +} + +#[derive(Identifiable, Queryable, Associations)] +#[belongs_to(Subdomain)] +#[table_name="urls"] +pub struct Url { + pub id: i64, + pub subdomain_id: i64, + pub status: u16, + pub body: Vec, +} diff --git a/src/schema.rs b/src/schema.rs new file mode 100644 index 00000000..6dcfec57 --- /dev/null +++ b/src/schema.rs @@ -0,0 +1,53 @@ +table! { + domains (id) { + id -> Nullable, + value -> Text, + } +} + +table! { + ipaddrs (id) { + id -> Nullable, + family -> Text, + value -> Text, + } +} + +table! { + subdomain_ipaddrs (id) { + id -> Nullable, + subdomain_id -> Integer, + ip_addr_id -> Integer, + } +} + +table! { + subdomains (id) { + id -> Nullable, + domain_id -> Integer, + value -> Text, + } +} + +table! { + urls (id) { + id -> Nullable, + subdomain_id -> Integer, + value -> Text, + status -> Nullable, + body -> Nullable, + } +} + +joinable!(subdomain_ipaddrs -> domains (subdomain_id)); +joinable!(subdomain_ipaddrs -> ipaddrs (ip_addr_id)); +joinable!(subdomains -> domains (domain_id)); +joinable!(urls -> domains (subdomain_id)); + +allow_tables_to_appear_in_same_query!( + domains, + ipaddrs, + subdomain_ipaddrs, + subdomains, + urls, +); diff --git a/src/shell.rs b/src/shell.rs index 73eb862c..d53d5116 100644 --- a/src/shell.rs +++ b/src/shell.rs @@ -1,4 +1,8 @@ +use errors::*; + use colored::Colorize; +use diesel::prelude::*; +use migrations; // use rand::prelude::*; use rustyline; use rustyline::completion::Completer; @@ -23,6 +27,7 @@ pub enum Command { Run, Set, Show, + Update, Use(Vec), Interrupt, @@ -36,6 +41,7 @@ impl Command { Command::Run => "run", Command::Set => "set", Command::Show => "show", + Command::Update => "update", Command::Use(_) => "use", Command::Interrupt => unreachable!(), } @@ -49,6 +55,7 @@ impl Command { Command::Run.as_str(), Command::Set.as_str(), Command::Show.as_str(), + Command::Update.as_str(), Command::Use(Vec::new()).as_str(), ]; } @@ -174,6 +181,7 @@ impl Readline { "run" => Some(Command::Run), "set" => Some(Command::Set), "show" => Some(Command::Show), + "update" => Some(Command::Update), "use" => Some(Command::Use(cmd[1..].to_vec())), x => { eprintln!("Error: unknown command: {:?}", x); @@ -212,8 +220,7 @@ pub fn print_banner() { "#, "osint".green(), "recon".green(), "security".green()); } - -pub fn run() -> () { +pub fn run() -> Result<()> { print_banner(); // wait("checking tor status"); @@ -221,13 +228,21 @@ pub fn run() -> () { // println!("\x1b[1m[\x1b[32m*\x1b[0;1m]\x1b[0m updating registry..."); // wait("updating registry"); + let _db = worker::spawn_fn("Connecting to database", || { + let db = SqliteConnection::establish("foo.db") + .context("Failed to connect to database")?; + migrations::run(&db) + .context("Failed to run migrations")?; + Ok(db) + }, false)?; + let mut rl = Readline::new(); loop { match rl.next() { Some(Command::Add) => println!("add"), Some(Command::Back) => if rl.take_module().is_none() { - break + break; }, Some(Command::Run) => { if let Some(module) = rl.module() { @@ -241,6 +256,11 @@ pub fn run() -> () { // TODO: set jobs 25 Some(Command::Set) => println!("set"), Some(Command::Show) => println!("show"), + Some(Command::Update) => { + // TODO + worker::spawn("Updating public suffix list"); + worker::spawn("Updating modules"); + }, Some(Command::Use(mut args)) => { if !args.is_empty() { rl.set_module(args.remove(0)); @@ -252,4 +272,6 @@ pub fn run() -> () { None => (), } } + + Ok(()) } diff --git a/src/worker.rs b/src/worker.rs index 3fbcbaaf..8f02ccab 100644 --- a/src/worker.rs +++ b/src/worker.rs @@ -1,5 +1,7 @@ +use errors::*; + use std::time::Duration; -use std::sync::mpsc; +use std::sync::{mpsc, Arc, Mutex}; use std::thread; use term::Spinner; @@ -36,3 +38,40 @@ pub fn spawn(task: &str) { // spinner.clear(); spinner.done(); } + +pub fn spawn_fn(label: &str, f: F, clear: bool) -> Result + where F: FnOnce() -> Result { + let (tx, rx) = mpsc::channel(); + + let spinner = Arc::new(Mutex::new(Spinner::random(label.to_string()))); + let spinner2 = spinner.clone(); + let t = thread::spawn(move || { + let mut spinner = spinner2.lock().unwrap(); + + loop { + match rx.recv_timeout(Duration::from_millis(100)) { + Ok(Event::Info(info)) => spinner.log(&info), + Ok(Event::Done) => break, + Err(mpsc::RecvTimeoutError::Timeout) => (), + Err(mpsc::RecvTimeoutError::Disconnected) => break, + } + spinner.tick(); + } + }); + + // run work in main thread + let result = f()?; + tx.send(Event::Done).unwrap(); + + t.join().expect("thread failed"); + + let spinner = spinner.lock().unwrap(); + + if clear { + spinner.clear(); + } else { + spinner.done(); + } + + Ok(result) +}