#define _CRT_SECURE_NO_WARNINGS #include <iostream> #include <string> #include <map> #include <vector> #include <time.h> #include "sqlite3.h"
#include <boost/format.hpp> #include <boost/tokenizer.hpp> #include <boost/lexical_cast.hpp>
using namespace std; using namespace boost;
sqlite3* open_database(std::string database_name) { int ref = -1; sqlite3 *db = 0;
ref = sqlite3_open(database_name.c_str(), &db); if (ref == SQLITE_OK) return db; return false; } bool close_database(sqlite3 *db) { int ref = sqlite3_close(db); if (ref == SQLITE_OK) return true; return false; } bool exec_sql(sqlite3 *db, char *sql) { char *error_code = 0; int ref = sqlite3_exec(db, sql, 0, 0, &error_code); if (ref == SQLITE_OK) { return true; } return false; }
void Init_Database() { sqlite3* open_db = open_database("./database.db"); if (open_db != false) { std::string sql = "create table HostDB(" "uid primary key," "host_address char(128) not null," "host_username char(128) not null," "host_password char(128) not null," "host_port char(128) not null," "host_group char(128) not null default 'DefaultGroup'" ");"; char run_sql[1024] = { 0 }; strcpy(run_sql, sql.c_str()); exec_sql(open_db, run_sql); } close_database(open_db); }
void AddHost_DB(sqlite3* open_db, std::string address, std::string username, std::string password, std::string port) { std::string format_string = boost::str(boost::format("insert into HostDB(host_address,host_username,host_password,host_port) values('%s','%s','%s','%s');") % address %username %password %port); char run_sql[2048] = { 0 }; strcpy(run_sql, format_string.c_str()); bool ref = exec_sql(open_db, run_sql); if (ref == true) { std::cout << "[+] 增加主机: " << address << " 完成" << std::endl; } }
void DeleteHost_DB(sqlite3 *open_db, std::string address) { std::string format_string = boost::str(boost::format("delete from HostDB where host_address = '%s';") % address); char run_sql[2048] = { 0 }; strcpy(run_sql, format_string.c_str()); bool ref = exec_sql(open_db, run_sql); if (ref == true) { std::cout << "[-] 删除主机: " << address << " 完成" << std::endl; } }
void SetHostGroup_DB(sqlite3* open_db, std::string address, std::string group_name) { std::string format_string = boost::str(boost::format("update HostDB set host_group='%s' where host_address = '%s';") % group_name %address); char run_sql[2048] = { 0 }; strcpy(run_sql, format_string.c_str()); bool ref = exec_sql(open_db, run_sql); if (ref == true) { std::cout << "[+] 主机: " << address << " 已加入到: " << group_name << " 组" << std::endl; } }
void ShowHostGroup_DB(sqlite3 *open_db) { sqlite3_stmt *stmt = 0; std::string format_string = "SELECT host_group,count(*) FROM HostDB GROUP BY host_group;";
char run_sql[1024] = { 0 }; strcpy(run_sql, format_string.c_str());
int ref = sqlite3_prepare_v2(open_db, run_sql, -1, &stmt, 0); if (ref == SQLITE_OK) { while (sqlite3_step(stmt) == SQLITE_ROW) { const unsigned char *host_group = sqlite3_column_text(stmt, 0); int host_group_count = sqlite3_column_int(stmt, 1);
std::cout << host_group << " " << host_group_count << std::endl; } } sqlite3_finalize(stmt); }
void ShowHost_DB(sqlite3 *open_db) { sqlite3_stmt *stmt = 0; std::string format_string = "select * from HostDB;";
char run_sql[1024] = { 0 }; strcpy(run_sql, format_string.c_str());
int ref = sqlite3_prepare_v2(open_db, run_sql, -1, &stmt, 0); if (ref == SQLITE_OK) { while (sqlite3_step(stmt) == SQLITE_ROW) { const unsigned char *host_address = sqlite3_column_text(stmt, 1); const unsigned char *host_username = sqlite3_column_text(stmt, 2); const unsigned char *host_paddword = sqlite3_column_text(stmt, 3); const unsigned char *host_port = sqlite3_column_text(stmt, 4); const unsigned char *host_group = sqlite3_column_text(stmt, 5);
std::cout << host_address << " " << host_username << " " << host_paddword << " " << host_port << " " << host_group << std::endl; } } sqlite3_finalize(stmt); }
void ShowGroupHostList(sqlite3 *open_db, std::string group_name) { sqlite3_stmt *stmt = 0; std::string format_string = boost::str(boost::format("select * from HostDB where host_group = '%s';") % group_name); char run_sql[2048] = { 0 }; strcpy(run_sql, format_string.c_str());
int ref = sqlite3_prepare_v2(open_db, run_sql, -1, &stmt, 0); if (ref == SQLITE_OK) { std::cout << "----------------------------------------------------------" << std::endl; std::cout << "主机组: " << group_name << std::endl; std::cout << "----------------------------------------------------------" << std::endl; while (sqlite3_step(stmt) == SQLITE_ROW) { const unsigned char *host_address = sqlite3_column_text(stmt, 1); const unsigned char *host_username = sqlite3_column_text(stmt, 2); const unsigned char *host_port = sqlite3_column_text(stmt, 4);
std::cout << host_address << " " << host_username << " " << host_port << std::endl; } } sqlite3_finalize(stmt); }
int main(int argc, char const *argv[]) { sqlite3* open_db = open_database("./database.db"); Init_Database(); std::string command;
while (1) { std::cout << "[ LyShark Shell ] # "; std::getline(std::cin, command);
if (command.length() == 0) { continue; } else if (command == "help") { std::cout << "帮助菜单" << std::endl; } else { boost::char_separator<char> sep(", --"); typedef boost::tokenizer<boost::char_separator<char>> CustonTokenizer; CustonTokenizer tok(command, sep); std::vector<std::string> vecSegTag; for (CustonTokenizer::iterator beg = tok.begin(); beg != tok.end(); ++beg) { vecSegTag.push_back(*beg); } if (vecSegTag.size() == 9 && vecSegTag[0] == "AddHost") { if (vecSegTag[1] == "address" && vecSegTag[3] == "username" && vecSegTag[5] == "password" && vecSegTag[7] == "port") { std::string set_address = vecSegTag[2]; std::string set_username = vecSegTag[4]; std::string set_password = vecSegTag[6]; std::string set_port = vecSegTag[8]; AddHost_DB(open_db, set_address, set_username, set_password, set_port); } } else if (vecSegTag.size() == 3 && vecSegTag[0] == "DeleteHost") { if (vecSegTag[1] == "address") { std::string set_address = vecSegTag[2]; DeleteHost_DB(open_db, set_address); } } else if (vecSegTag.size() == 5 && vecSegTag[0] == "SetHostGroup") { if (vecSegTag[1] == "address" && vecSegTag[3] == "group") { std::string set_address = vecSegTag[2]; std::string set_group = vecSegTag[4]; SetHostGroup_DB(open_db, set_address, set_group); } } else if (vecSegTag.size() == 1 && vecSegTag[0] == "ShowHost") { std::cout << "-----------------------------------------------------------------------------" << std::endl; std::cout << "IP地址 " << "用户名 " << "密码 " << "端口号 " << "默认组 " << std::endl; std::cout << "-----------------------------------------------------------------------------" << std::endl; ShowHost_DB(open_db); } else if (vecSegTag.size() == 1 && vecSegTag[0] == "ShowHostGroup") { std::cout << "-----------------------------------------------------------------------------" << std::endl; std::cout << "主机组名 " << "主机数量 " << std::endl; std::cout << "-----------------------------------------------------------------------------" << std::endl; ShowHostGroup_DB(open_db); } else if (vecSegTag.size() == 3 && vecSegTag[0] == "ShowGroupHostList") { if (vecSegTag[1] == "group") { std::string set_group = vecSegTag[2]; ShowGroupHostList(open_db, set_group); } } } }
close_database(open_db); return 0; }
|