Edited at
DiscordDay 24

ウイルススキャンbotを作る


ウイルススキャンbotを作る

なんとなく思いついたので作ります

VirusTotalというサイトがAPI公開してるのでそれ使います。

ただ、1分に4回という制限があるのでサーバーによっては使いものにならないかもしれないです。


どんなの


  • ファイルがアップロードされたらそのファイルのスキャンリクエストを投げる

  • スキャンが終わったらそのファイルの結果をファイルが送られたチャンネルに送信する。

注意点として、VirusTotal Intelligenceを使うとスキャンしたファイルを閲覧することができるため他人に見られたら困るようなファイルは上げないでください。


あったら良さそうだけど現実的じゃないやつ


  1. ファイルがアップロードされたらその瞬間にメッセージを編集してスキャン中の旨を表示する。

  2. URLもスキャンする

なぜ現実的ではないかと言うとAPI制限があるから

1に関しては 1分間に5個以上のファイルが投稿された場合待たないとファイルがダウンロードできない

2はファイル以上に投稿頻度が高いためスキャンが追い付かない


作ってく

言語はElixirで作ります。

以下のライブラリを使います

- discord_alchemy

- virus_total

今回はnanobox上で動かします

nanoboxくんすこってあげてください

-> 現在パッケージとか更新されてないので使わない方が良さげ...

適当なディレクトリに入ってboxfileを作っていきます

もし音声とか使いたいならここでffmpeg入れるようにします。

今回はテキストだけなので何も書きませんでした。


boxfile.yml

run.config:

engine: elixir

engine.config:
runtime: elixir-1.5.2 #最新のやつ


作ったら

$ nanobox run

これでelixirのコンテナが作られて自動的に接続されるのでプロジェクトを作っていきましょう

マウントされてるディレクトリは/appになるのでcd /で移動して以下のコマンドでプロジェクトを作ります

$ mix new app --app virus_scanner

既にディレクトリが存在するけど?みたいなこと言われると思いますがYesで行きましょう

Mixfileを以下のようにします


mix.exs

  # Run "mix help deps" to learn about dependencies.

defp deps do
[
# {:dep_from_hexpm, "~> 0.3.0"},
# {:dep_from_git, git: "https://github.com/elixir-lang/my_dep.git", tag: "0.1.0"},
{:alchemy, "~> 0.6.1", hex: :discord_alchemy},
{:virus_total, "~> 0.0.2"},
{:httpoison, "~> 0.13.0", override: true}
]
end

{:httpoison, "~> 0.13.0", override: true}に関してはalchemy(0.13.0)とvirus_total(0.9.0)で競合するのでoverride: trueにして0.13.0を使うようにしてます。

alchemyのドキュメントを読むとmix alchemy.initしろみたいに書いてありますが、mybotとかいうのが作られるだけなので無視して自分で書いていきます。

Discord BotとVirusTotalのトークンは適当に作って用意しておいてください

https://github.com/reactiflux/discord-irc/wiki/Creating-a-discord-bot-&-getting-a-token

ではトークンをconfig.exsに記載しましょう。


config/config.exs

......

# It is also possible to import configuration files, relative to this
# directory. For example, you can emulate configuration per environment
# by uncommenting the line below and defining dev.exs, test.exs and such.
# Configuration from the imported file will override the ones defined
# here (which is why it is important to import them last).
#
# import_config "#{Mix.env}.exs"

config :virus_scanner,
discord_token: "とーくん",
virus_total_token: "とーくん"


次にこのトークンを使いBotを動かすためのコードを書きます

とりあえず、以下の様にしました


lib/virus_scanner.ex

defmodule VirusScanner do

use Application
alias Alchemy.Client

@token Application.get_env(:virus_scanner, :discord_token)

def start(_type, _args) do
Client.start(@token)
end
end


適当なサーバーにBotを追加してmix run --no-haltと実行すれば動くはずです。

このままだと、存在してるだけのBotなのでイベントハンドラを追加します。

サンプルにあるような発言したメッセージを出力するものを書きたいと思います

以下の様にしてみました


lib/events/on_message.ex

defmodule VirusScanner.OnMessage do

use Alchemy.Events

Events.on_message(:inspect)

def inspect(message) do
IO.inspect message.content
end
end



lib/virus_scanner.ex

.....

def start(_type, _args) do
run = Client.start(@token)
use VirusScanner.OnMessage
run
end


これで再度実行してBotが入ってるサーバーで発言してみましょう

image.png

無事出力されています。

ここから、もしファイルだったらVirusTotalにスキャンリクエストを送るという挙動にしていきます。

ファイルだった場合%Alchemy.Messageのattachmentsに%Alchemy.Attachmentが入っているためそこの確認をして分岐したいと思います


lib/events/on_message.ex

defmodule VirusScanner.OnMessage do

use Alchemy.Events

Events.on_message(:scan)

def scan(%Alchemy.Message{attachments: [%Alchemy.Attachment{url: url}]}) do
IO.inspect url
end

def scan(_msg), do: IO.inspect "on message"
end


こんな感じにすればいい感じにパターンマッチでうごくと思います。

image.png

良い感じです

ただ、これだと写真もスキャン対象になってしまいます。

なので正規表現を使って画像以外のスキャンをするように、しましょう。

ついでにトークンも追加


lib/events/on_message.ex

defmodule VirusScanner.OnMessage do

use Alchemy.Events

@virusTotalToken Application.get_env(:virus_scanner, :virus_total_token)

Events.on_message(:scan)

def scan(%Alchemy.Message{attachments: [%Alchemy.Attachment{url: url}]}) do
regex = Regex.compile!("/\.gif$|\.png$|\.jpg$|\.jpeg$|\.bmp$/i")

unless Regex.match?(regex, url) do
IO.inspect url
end
end

def scan(_msg), do: IO.inspect "on message"
end


こんな感じになりました。

これをVirusTotalのURLスキャンを使ってスキャンするようにします。

まず、スキャンしたファイルの名前やスキャンIDなどを記録しておく物が必要になります。

今回はErlang Term Storage(ETS)を使います。

すごく雑に言えばRDBみたいなものです。

ではETSを使うためのコードを書いていきます。


lib/virus_scanner.ex

defmodule VirusScanner do

use Application
alias Alchemy.Client

@token Application.get_env(:virus_scanner, :discord_token)

def start(_type, _args) do
init_storage()
run = Client.start(@token)
use VirusScanner.OnMessage
run
end

def init_storage() do
:ets.new(:virus_scanner, [:named_table, :public])
end
end


こんな感じにnamed_tableを使ってテーブルを作ります。

そしたら、先ほどの正規表現でいい感じ(たぶん)なやつにスキャンとinsertするコード書いていきます。


lib/events/on_message.ex

defmodule VirusScanner.OnMessage do

use Alchemy.Events

import Alchemy.Embed

alias Alchemy.Embed
alias Alchemy.User

@virusTotalToken Application.get_env(:virus_scanner, :virus_total_token)

Events.on_message(:scan)

def scan(%Alchemy.Message{attachments: [%Alchemy.Attachment{url: url, filename: filename}], author: author, channel_id: channel_id}) do
regex = Regex.compile!("/\.gif$|\.png$|\.jpg$|\.jpeg$|\.bmp$/i")

message = %{
channel_id: channel_id
}

unless Regex.match?(regex, url) do
{_, result} = VirusTotal.Client.url_scan(@virusTotalToken, url)
result = Enum.into(result, %{})

if Map.get(result, "response_code") == 1 do
:ets.insert(:virus_scanner, {Map.get(result, "scan_id"), filename})
end

%Embed{}
|> title("スキャンリクエストを送りました")
|> description("ファイル名: #{filename}")
|> author(name: author.username, icon_url: User.avatar_url(author))
|> field("レスポンスコード", Map.get(result, "response_code"))
|> field("レスポンスメッセージ", Map.get(result, "verbose_msg"))
|> color(0x00ff00)
|> Embed.send
end
end

def scan(_msg), do: nil
end


とりあえずこんな感じになりました。

messageにchannel_idを入れてるのはEmbedのマクロで必要だからです。

あとEmbed使ってるのでおしゃんてい :blush:

image.png

その後response_codeが1(成功)だったらetsにinsertみたいな感じです。

エラーだった時メッセージを変えるみたいなことしたかったのですが、12/24当日に書いてるので時間の都合上できませんでした :bow:

定期実行するためにquantumというライブラリが必要だったので適当に追加します

config.exsに以下を追記


config/config.exs

config :porcelain, driver: Porcelain.Driver.Basic

config :quantum,
timezone: "Asia/Tokyo"

config :virus_scanner, VirusScanner.Scheduler,
jobs: [
{"* * * * *", fn -> VirusScanner.OnMessage.check_reports() end}
]


scheduler.exを追加


lib/scheduler.ex

defmodule VirusScanner.Scheduler do

use Quantum.Scheduler,
otp_app: :virus_scanner
end

virus_scanner.exにスーパーバイザを動かすコード書きます


lib/virus_scanner.ex

defmodule VirusScanner do

use Application
alias Alchemy.Client

@token Application.get_env(:virus_scanner, :discord_token)

def start(_type, _args) do
import Supervisor.Spec, warn: false

children = [
worker(VirusScanner, []),
worker(VirusScanner.Scheduler, [])
]

opts = [strategy: :one_for_one, name: YourApp.Supervisor]

Supervisor.start_link(children, opts)
end

def start_link() do
init_storage()
run = Client.start(@token)
use VirusScanner.OnMessage
run
end

def init_storage() do
:ets.new(:virus_scanner, [:named_table, :public, :set])
end
end


たぶん、こんな感じでいいです。

その後config.exsに書いたVirusScanner.OnMessage.check_reportsを定義します

本来別のモジュール作ってやるべきだと思いますが時間の都合上on_message.exに全部書いていきます


lib/events/on_message.ex

defmodule VirusScanner.OnMessage do

use Alchemy.Events

import Alchemy.Embed

alias Alchemy.Embed
alias Alchemy.User

@virusTotalToken Application.get_env(:virus_scanner, :virus_total_token)

Events.on_message(:scan)

def scan(%Alchemy.Message{attachments: [%Alchemy.Attachment{url: url, filename: filename}], author: author, channel_id: channel_id}) do
regex = Regex.compile!("/\.gif$|\.png$|\.jpg$|\.jpeg$|\.bmp$/i")

message = %{
channel_id: channel_id
}

unless Regex.match?(regex, url) do
{_, result} = VirusTotal.Client.url_scan(@virusTotalToken, url)
result = Enum.into(result, %{})

if Map.get(result, "response_code") == 1 do
:ets.insert(:virus_scanner, {Map.get(result, "scan_id"), filename, channel_id})
end

%Embed{}
|> title("スキャンリクエストを送りました")
|> description("ファイル名: #{filename}")
|> author(name: author.username, icon_url: User.avatar_url(author))
|> field("レスポンスコード", Map.get(result, "response_code"))
|> field("レスポンスメッセージ", Map.get(result, "verbose_msg"))
|> color(0x00ff00)
|> Embed.send
end
end

def scan(_msg), do: nil

def check_reports() do
:ets.match_object(:virus_scanner, {:"$1", :"_", :"_"})
|> Enum.map(fn(x) -> get_report(x) end)
end

def get_report({scan_id, filename, channel_id}) do
message = %{
channel_id: channel_id
}

{:ok, report} = VirusTotal.Client.url_report(@virusTotalToken, scan_id)
report = Enum.into(report, %{})

if Map.get(report, "response_code") == 1 do
%Embed{}
|> title("スキャンが完了しました")
|> description("ファイル名: #{filename}")
|> field("レスポンスコード", Map.get(report, "response_code"), inline: true)
|> field("検出数", Map.get(report, "positives"), inline: true)
|> field("レスポンスメッセージ", Map.get(report, "verbose_msg"))
|> color(0xff0000)
|> Embed.send

:ets.delete(:virus_scanner, scan_id)
end
end
end


こんな感じになりました

これをmix run --no-haltで起動してBotが入ってるサーバーに適当なファイルを上げてみてください。

image.png

いいいいかんじです :smile:


最後に

時間が無かったのでかなり適当に作りましたが割といい感じに動いたのでよかったです。

ソースは気が向いたらgithubに上げておきます。

では、クリスマスイブにプログラミンしてるぼくはケーキ食ってねます :joy: