6.14 筆記法實例 (commands.py)

import os
import getopt

from viper.common.out import *
from viper.common.colors import bold, cyan, white
from viper.core.session import __session__
from viper.core.plugins import __modules__
from viper.core.database import Database
from viper.core.storage import store_sample, get_sample_path

class Commands(object):

    def __init__(self):
        # Open connection to the database.
        self.db = Database()

        # Map commands to their related functions.
        self.commands = dict(
            help=dict(obj=self.cmd_help, description="Show this help message"),
            open=dict(obj=self.cmd_open, description="Open a file"),
            close=dict(obj=self.cmd_close, description="Close the current session"),
            info=dict(obj=self.cmd_info, description="Show information on the opened file"),
            clear=dict(obj=self.cmd_clear, description="Clear the console"),
            store=dict(obj=self.cmd_store, description="Store the opened file to the local repository"),
            delete=dict(obj=self.cmd_delete, description="Delete the opened file"),
            find=dict(obj=self.cmd_find, description="Find a file"),
        )

    ##
    # CLEAR
    #
    # This command simply clears the shell.
    def cmd_clear(self, *args):
        os.system('clear')

    ##
    # HELP
    #
    # This command simply prints the help message.
    # It lists both embedded commands and loaded modules.
    def cmd_help(self, *args):
        print(bold("Commands:"))

        rows = []
        for command_name, command_item in self.commands.items():
            rows.append([command_name, command_item['description']])

        print(table(['Command', 'Description'], rows))       
        print("")
        print(bold("Modules:"))

        rows = []
        for module_name, module_item in __modules__.items():
            rows.append([module_name, module_item['description']])

        print(table(['Command', 'Description'], rows))

    ##
    # OPEN
    #
    # This command is used to open a session on a given file.
    # It either can be an external file path, or a SHA256 hash of a file which
    # has been previously imported and stored.
    # While the session is active, every operation and module executed will be
    # run against the file specified.
    def cmd_open(self, *args):
        def usage():
            print("usage: open [-h] [-f] target")

        def help():
            usage()
            print("")
            print("Options:")
            print("\t--help (-h)\tShow this help message")
            print("\t--file (-f)\tThe target is a file")
            print("")

        try:
            opts, argv = getopt.getopt(args, 'hf', ['help', 'file'])
        except getopt.GetoptError as e:
            print(e)
            usage()
            return

        is_file = False

        for opt, value in opts:
            if opt in ('-h', '--help'):
                help()
                return
            elif opt in ('-f', '--file'):
                is_file = True

        if len(argv) == 0:
            usage()
            return
        else:
            target = argv[0]

        if is_file:
            target = os.path.expanduser(target)

            if not os.path.exists(target) or not os.path.isfile(target):
                print_error("File not found")
                return

            __session__.set(target)
        else:
            target = argv[0].strip().lower()
            path = get_sample_path(target)
            if path:
                __session__.set(path)

    ##
    # CLOSE
    #
    # This command resets the open session.
    # After that, all handles to the opened file should be closed and the
    # shell should be restored to the default prompt.
    def cmd_close(self, *args):
        __session__.clear()

    ##
    # INFO
    #
    # This command returns information on the open session. It returns details
    # on the file (e.g. hashes) and other information that might available from
    # the database.
    def cmd_info(self, *args):
        if __session__.is_set():
            print(table(
                ['Key', 'Value'],
                [
                    ('Name', __session__.file.name),
                    ('Path', __session__.file.path),
                    ('Size', __session__.file.size),
                    ('Type', __session__.file.type),
                    ('MD5', __session__.file.md5),
                    ('SHA1', __session__.file.sha1),
                    ('SHA256', __session__.file.sha256),
                    ('SHA512', __session__.file.sha512),
                    ('SSdeep', __session__.file.ssdeep),
                    ('CRC32', __session__.file.crc32)
                ]
            ))

    ##
    # STORE
    #
    # This command stores the opened file in the local repository and tries
    # to store details in the database.
    def cmd_store(self, *args):
        # TODO: Add tags argument.
        if __session__.is_set():
            # Store file to the local repository.
            new_path = store_sample(__session__.file)
            # Add file to the database.
            status = self.db.add(__session__.file)

            print_success("Stored to: {0}".format(new_path))

            # Open session to the new file.
            self.cmd_open(*[__session__.file.sha256])

    ##
    # DELETE
    #
    # This commands deletes the currenlty opened file (only if it's stored in
    # the local repository) and removes the details from the database
    def cmd_delete(self, *args):
        if __session__.is_set():
            while True:
                choice = raw_input("Are you sure you want to delete this binary? Can't be reverted! [y/n] ")
                if choice == 'y':
                    break
                elif choice == 'n':
                    return

            rows = self.db.find('sha256', __session__.file.sha256)
            if rows:
                malware_id = rows[0].id
                if self.db.delete(malware_id):
                    print_success("File deleted")
                else:
                    print_error("Unable to delete file")

            os.remove(get_sample_path(__session__.file.sha256))
            __session__.clear()

    ##
    # FIND
    #
    # This command is used to search for files in the database.
    def cmd_find(self, *args):
        if len(args) == 0:
            print_error("Invalid search term")
            return

        key = args[0]
        try:
            value = args[1]
        except IndexError:
            value = None

        items = self.db.find(key, value)
        if not items:
            return

        rows = []
        for item in items:
            rows.append([item.name, item.type, item.size, item.sha256])

        print(table(['Name', 'Type', 'Size', 'SHA256'], rows))

此程式主要功能為定義 viper 系統指令。

有 1 類別 Commands
及 8 類別函式
    cmd_clear, 
    cmd_close, 
    cmd_delete, 
    cmd_find, 
    cmd_help, 
    cmd_info, 
    cmd_open, 
    cmd_store 

類別 instance 初始化時，做兩件事
1. 用 self.db = Database() 開啟與 db 的連線

2. 定義 self.commands
self.commands 是用兩層 巢狀字典組成
字典第一層 key 值是系統指令， value 是第二層字典。
第二層字典內有兩個 item。
    第一個 item 是 obj，對應到類別函式，
    第二個 item 是 description，即描述指令功能之文字

因此 console.py 只要呼叫 self.cmd.commands[第一層(指令)][第二層(obj)]()
就可執行 viper 系統函式。

接下來逐一看類別函式。

1. cmd_clear()
此函式用 os.system('clear') 清除 shell 上訊息。

2. cmd_help()
此函式會印出 help 訊息，即系統指令與模組指令。
首先印出 Commands
並用迴圈 將 commands 中 items 抓出來，

所以 command_name 就是 key (指令), command_item 就是 value(obj 與 description)。
隨後將 key 與 value 塞進 list
再將此 list 塞進 rows (也是個 list)

接著用 prettytable 的 table 將系統指令以表格印出。

同樣的方法運用至將模組指令與描述以 table 印出。
(plugins.py 中， line 30 將模組指令以與系統指令相同之巢狀字典方式，組成)


3. cmd_open
此函式首先接收 *args，不定參數，也就是可以接參數，但參數數量不限。
函式中，定義了兩個函式 usage() 與 help()
usage() 印出 open 使用方法簡單說明
help() 中印出方法簡單說明與參數詳細說明

接者用 getopt 開始 parse 參數
getopt 回傳兩 item。
1. opts
2. argv

opts 是個 list of (option, value)
argv 是其他沒被定義到的參數值 list

用 try 來 parse 參數
若執行失敗，則 raise getopt.GetoptError 
接著再印出 usage()
再 return，結束 cmd_open

接下來設定 is_file變數，預設是 False

再來用迴圈處理 opts
主要將 option 拿出來，看到底是 -h 還是 -f
若是 -h，則印出 help()後直接 return 結束 cmd_open
若是 -f 則將 is_file 改為 True

接著變檢查 argv，也就是沒被定義的參數 list 長度是否為零
如果是零，則表示 -f 後面沒有接檔案路徑。
所以印出 usage()後，接 return 結束 cmd_open
若 argv 長度非零，則將 list 中第一個 item 抓為檔案位置，也就是 target。

接著，檢查 is_file flag 是否為 True，
若是，則用 os.path.expanduser() 處理 target 路徑，主要功能是將 ~ replace 成 $HOME。
接著用 os.path.exists() 檢查路徑是否存在，以及使用 os.path.isfile()檢查是否為檔案。
若不存在或不是檔案，則印出 File not found 錯誤訊息

若存在，則用 __session__.set() 開啟 session。

若 is_file flag 為 False (此種作法就是直接 open file)
則用 .strip().lower() 方法將 argv[0] 先以空白為區隔丟進 list，再將 list 裡面字串通通轉為小寫。
接著用 get_sample_path，取得檔案路徑
若路徑 (path) 不是空的
則用 __session__.set() 開啟 session。

4. cmd_close()
負責關掉 __session__

5. cmd_info()
負責印出現有 session 所分析檔案之基本資訊。
使用 prettytable 以表格漂亮印出相關資訊。

6. cmd_store()
將 session打開的檔案存在 local repository
將 檔案資訊存在 database。

首先用 __session__.is_set() 檢查 session 是否開啟。
接著用 store_sample(儲存檔案)後 ，將檔案路徑傳給 new_path 變數
再用 seld.db.add(__session__.file) 將檔案資訊存到 db 中

最後印出檔案儲存路徑，並重新開啟檔案 session。

7. cmd_delete()
首先檢查 session 是否開啟。
若開啟，則用一 while 迴圈，
再次確認使用者是否真的要刪除 current open file。
若是 則 break 迴圈，若否則用 return 結束 cmd_delete。

若是，則用 self.db.find('sha256', __session__.file.sha256) 找出檔案，並將結果 assign 給 rows
rows 的 type 是 list。
(Return the results represented by this Query as a list.)
參考文獻：http://docs.sqlalchemy.org/en/latest/orm/query.html?highlight=query

接著，若 rows 不是空，則用 rows[0].id 將 malware_id 找出來。
其中 rows[0] type 是 class instance。
並用 self.db.delete(malware.id) 刪除資料庫中該檔案資訊。
若失敗，則印出 Unable to delete file。

為何 rows[0] 一定會是 current opened file 呢？
原因在於 database.py 中 line 41 ~ line 49
針對 md5, crc32, sha1, sha256, sha512 欄位，做 uniq。

最後用 os.remove() 將 local repository 中檔案刪除
並關掉 session。

8. cmd_find()
搜尋資料庫中檔案資訊

首先檢查參數長度是否為零，若是，則表示沒給參數。
接著開始 parse 參數，key 表示 sha256, value 表示 hash value。
接著用 self.db.find(key, value) 將資料撈出來，倒給 items。
若無資料，就用 return 結束 cmd_find。
若有，則用 prettytable 印出檔案相關資訊。