#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Name: Trojan Manager Dev: K4YT3X Date Created: July 8, 2018 Last Modified: March 5, 2019 Licensed under the GNU General Public License Version 3 (GNU GPL v3), available at: https://www.gnu.org/licenses/gpl-3.0.txt Copyright (C) 2018-2019, K4YT3X <k4yt3x@k4yt3x.com> """ from avalon_framework import Avalon from prettytable import PrettyTable import hashlib import MySQLdb import readline import sys import traceback VERSION = '1.3.7' COMMANDS = [ "CreateUserTable", "TruncateUserTable", "Verify", "AddUser", "DelUser", "Show", "SetQuota", "AddQuota", "ClearUsage", "Exit", "Quit", ] def show_affection(function): """ Shows cursor execution affected rows """ def wrapper(*args, **kwargs): function(*args, **kwargs) Avalon.debug_info('{} row(s) affected'.format(args[0].cursor.rowcount)) return wrapper def catch_mysql_errors(function): """ Catch mysqldb warnings and errors """ def wrapper(*args, **kwargs): try: function(*args, **kwargs) except (MySQLdb.Error, MySQLdb.Warning) as e: Avalon.error(e) return 1 return wrapper class ShellCompleter(object): """ Completer for readline This class tries to match input text to a list of available commands. """ def __init__(self, options): self.options = sorted(options) def complete(self, text, state): if state == 0: if text: self.matches = [s for s in self.options if s and s.lower().startswith(text.lower())] else: self.matches = self.options[:] try: return self.matches[state] except IndexError: return None class TrojanDatabase: """ Trojan Database Connector This class generates objects that connects to and controls the trojan database. """ def __init__(self, db_host, db_user, db_pass, db, table): """ Initialize database connection """ self.db_host = db_host self.db_user = db_user self.db_pass = db_pass self.db = db self.table = table self.connection = MySQLdb.connect(self.db_host, self.db_user, self.db_pass, self.db) self.cursor = self.connection.cursor() @show_affection @catch_mysql_errors def create_user_table(self): """ Create new user table """ sql = ["CREATE TABLE {} (".format(self.table), "id INT UNSIGNED NOT NULL AUTO_INCREMENT,", "username VARCHAR(64) NOT NULL,", "password CHAR(56) NOT NULL,", "quota BIGINT NOT NULL DEFAULT 0,", "download BIGINT UNSIGNED NOT NULL DEFAULT 0,", "upload BIGINT UNSIGNED NOT NULL DEFAULT 0,", "PRIMARY KEY (id),", "INDEX (password)", ");", ] self.cursor.execute(''.join(sql)) self.connection.commit() return 0 @show_affection @catch_mysql_errors def truncate_user_table(self): """ truncate the user table """ self.cursor.execute('TRUNCATE {};'.format(self.table)) self.connection.commit() return 0 @show_affection @catch_mysql_errors def drop_user_table(self): """ Boom """ self.cursor.execute('DROP TABLE {};'.format(self.table)) self.connection.commit() return 0 @show_affection @catch_mysql_errors def add_user(self, username, password): """ Add new user into database """ self.cursor.execute("SELECT * FROM {} WHERE username = '{}'".format(self.table, username)) if self.cursor.fetchone() is not None: Avalon.error('User {} already exists'.format(username)) self.cursor.rowcount = 0 # No actual changes to database return 1 fullhash = hashlib.sha224('{}:{}'.format(username, password).encode('utf-8')).hexdigest() self.cursor.execute("INSERT INTO {} (username, password) VALUES ('{}', '{}')".format(self.table, username, fullhash)) # self.cursor.execute("INSERT INTO {} SHA2(CONCAT(username, ':', password)', 224) VALUES ('{}', '{}')".format(self.table, username, password)) # self.connection.commit() return 0 @show_affection @catch_mysql_errors def del_user(self, username): """ Delete a user from the database """ self.cursor.execute("DELETE FROM {} WHERE username = '{}'".format(self.table, username)) self.connection.commit() return 0 @catch_mysql_errors def user_exists(self, username): """ Determines if a user exists in the database """ self.cursor.execute("SELECT * FROM {} WHERE username = '{}'".format(self.table, username)) user_id = self.cursor.fetchone() if user_id is not None: return True return False @catch_mysql_errors def show_users(self, show_quota=False): """ List users If show_quota is True, include quota and data usage """ total_users = self.cursor.execute("SELECT * FROM {}".format(self.table)) columns = ['ID', 'Username', 'Password'] if show_quota: columns += ['Quota', 'Download', 'Upload'] table = PrettyTable(columns) for user in self.cursor.fetchall(): if show_quota: table.add_row(user) else: table.add_row([user[0], user[1], user[2]]) print(table) Avalon.info('Query complete, {} user(s) found in database'.format(total_users)) return 0 def convert_units(self, data): """ Convert data unit into bytes """ try: try: return int(data) except ValueError: if data[-1].lower() == 'k': return int(data[:-1]) * 1024 elif data[-1].lower() == 'm': return int(data[:-1]) * 1024 ** 2 elif data[-1].lower() == 'g': return int(data[:-1]) * 1024 ** 3 elif data[-1].lower() == 't': return int(data[:-1]) * 1024 ** 4 elif data[-1].lower() == 'p': return int(data[:-1]) * 1024 ** 5 else: return False except ValueError: return False @show_affection @catch_mysql_errors def set_quota(self, username, quota): """ Set user quota to a specific value """ converted = self.convert_units(quota) if not converted: Avalon.error('Invalid quota input') return 1 self.cursor.execute("UPDATE {} SET quota = {} WHERE username = '{}'".format(self.table, converted, username)) self.connection.commit() return 0 @show_affection @catch_mysql_errors def add_quota(self, username, appended_quota): """ Append quota """ converted = self.convert_units(appended_quota) if not converted: Avalon.error('Invalid quota input') return 1 self.cursor.execute("UPDATE {} SET quota = quota + {} WHERE username = '{}'".format(self.table, converted, username)) self.connection.commit() return 0 @show_affection @catch_mysql_errors def clear_usage(self, username): """ Clear user data usage """ self.cursor.execute("UPDATE {} SET download = 0, upload = 0 WHERE username = '{}'".format(self.table, username)) return 0 @catch_mysql_errors def verify(self, fullhash): """ Verify if user credentials are valid """ valid_hashes = [] self.cursor.execute("SELECT * FROM {}".format(self.table)) all_users = self.cursor.fetchall() for user in all_users: valid_hashes.append(user[2]) if fullhash in valid_hashes: Avalon.info('Valid user') return 0 Avalon.warning('Invalid user') return 1 def print_legal_info(): print('Trojan Manager {}'.format(VERSION)) print('(C) 2018 K4YT3X') print('Licensed under GNU GPL v3') def print_help(): help_lines = [ "\n{}Commands are not case-sensitive{}".format(Avalon.FM.BD, Avalon.FM.RST), "CreateUserTable", "TruncateUserTable", "Verify [hash]", "AddUser [username] [password]", "DelUser [username]", "Show (users / quota)", "SetQuota [quota]", "AddQuota [quota]", "ClearUsage [username]", "Interactive / int" "Exit / Quit", "", ] for line in help_lines: print(line) def command_interpreter(db_connection, commands): """ Trojan shell command interpreter """ try: # Try to guess what the user is saying possibilities = [s for s in COMMANDS if s.lower().startswith(commands[1])] if len(possibilities) == 1: commands[1] = possibilities[0] if commands[1].replace(' ', '') == '': result = 0 elif commands[1].lower() == 'help': print_help() result = 0 elif commands[1].lower() == 'createusertable': result = db_connection.create_user_table() elif commands[1].lower() == 'truncateusertable': Avalon.warning('By truncating you will LOSE ALL USER DATA') if Avalon.ask('Are you sure you want to truncate?'): result = db_connection.truncate_user_table() else: Avalon.warning('Operation canceled') result = 0 elif commands[1].lower() == 'dropusertable': Avalon.warning('By dropping the table you will LOSE ALL USER DATA') if Avalon.ask('Are you sure you want to drop the table?'): result = db_connection.drop_user_table() else: Avalon.warning('Operation canceled') result = 0 elif commands[1].lower() == 'verify': result = db_connection.verify(commands[2]) elif commands[1].lower() == 'adduser': result = db_connection.add_user(commands[2], commands[3]) elif commands[1].lower() == 'deluser': result = db_connection.del_user(commands[2]) elif commands[1].lower() == 'show': if commands[2].lower() == 'users': result = db_connection.show_users() elif commands[2].lower() == 'quota': result = db_connection.show_users(show_quota=True) elif commands[1].lower() == 'setquota': result = db_connection.set_quota(commands[2], commands[3]) elif commands[1].lower() == 'addquota': result = db_connection.add_quota(commands[2], commands[3]) elif commands[1].lower() == 'clearusage': result = db_connection.clear_usage(commands[2]) elif commands[1].lower() == 'exit' or commands[1].lower() == 'quit': Avalon.warning('Exiting') exit(0) elif len(possibilities) > 0: Avalon.warning('Ambiguous command \"{}\"'.format(commands[1])) print('Use \"Help\" command to list available commands') result = 1 else: Avalon.error('Invalid command') print('Use \"Help\" command to list available commands') result = 1 return result except IndexError: Avalon.error('Invalid arguments') print('Use \"Help\" command to list available commands') result = 0 def main(): """ Trojan Manager main function This function can only be executed when this file is not being imported. """ # Create database controller connection try: trojan_db = TrojanDatabase('127.0.0.1', 'trojan', 'thisisthetrojandbpassword', 'trojan', 'users') except (MySQLdb.OperationalError) as e: Avalon.error('Error establishing connection to MySQL/MariaDB') Avalon.error('Please check your settings') traceback.print_exc() exit(1) # Begin command interpreting try: if sys.argv[1].lower() == 'interactive' or sys.argv[1].lower() == 'int': print_legal_info() # Set command completer completer = ShellCompleter(COMMANDS) readline.set_completer(completer.complete) readline.parse_and_bind('tab: complete') # Launch interactive trojan shell prompt = '{}[trojan]> {}'.format(Avalon.FM.BD, Avalon.FM.RST) while True: command_interpreter(trojan_db, [''] + input(prompt).split(' ')) else: # Return to shell with command return value exit(command_interpreter(trojan_db, sys.argv[0:])) except IndexError: Avalon.warning('No commands specified') exit(0) except (KeyboardInterrupt, EOFError): Avalon.warning('Exiting') exit(0) except Exception: Avalon.error('Exception caught') traceback.print_exc() exit(1) if __name__ == '__main__': main() else: Avalon.warning('This file cannot be imported')