ウイルススキャンbotを作る
なんとなく思いついたので作ります
VirusTotalというサイトがAPI公開してるのでそれ使います。
ただ、1分に4回という制限があるのでサーバーによっては使いものにならないかもしれないです。
どんなの
- ファイルがアップロードされたらそのファイルのスキャンリクエストを投げる
- スキャンが終わったらそのファイルの結果をファイルが送られたチャンネルに送信する。
注意点として、VirusTotal Intelligenceを使うとスキャンしたファイルを閲覧することができるため他人に見られたら困るようなファイルは上げないでください。
あったら良さそうだけど現実的じゃないやつ
- ファイルがアップロードされたらその瞬間にメッセージを編集してスキャン中の旨を表示する。
- URLもスキャンする
なぜ現実的ではないかと言うとAPI制限があるから
1に関しては 1分間に5個以上のファイルが投稿された場合待たないとファイルがダウンロードできない
2はファイル以上に投稿頻度が高いためスキャンが追い付かない
作ってく
言語はElixirで作ります。
以下のライブラリを使います
今回はnanobox上で動かします
nanoboxくんすこってあげてください
-> 現在パッケージとか更新されてないので使わない方が良さげ...
適当なディレクトリに入ってboxfileを作っていきます
もし音声とか使いたいならここでffmpeg入れるようにします。
今回はテキストだけなので何も書きませんでした。
run.config:
engine: elixir
engine.config:
runtime: elixir-1.5.2 #最新のやつ
作ったら
$ nanobox run
これでelixirのコンテナが作られて自動的に接続されるのでプロジェクトを作っていきましょう
マウントされてるディレクトリは/app
になるのでcd /
で移動して以下のコマンドでプロジェクトを作ります
$ mix new app --app virus_scanner
既にディレクトリが存在するけど?みたいなこと言われると思いますがYesで行きましょう
Mixfileを以下のようにします
# 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に記載しましょう。
......
# 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を動かすためのコードを書きます
とりあえず、以下の様にしました
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なのでイベントハンドラを追加します。
サンプルにあるような発言したメッセージを出力するものを書きたいと思います
以下の様にしてみました
defmodule VirusScanner.OnMessage do
use Alchemy.Events
Events.on_message(:inspect)
def inspect(message) do
IO.inspect message.content
end
end
.....
def start(_type, _args) do
run = Client.start(@token)
use VirusScanner.OnMessage
run
end
これで再度実行してBotが入ってるサーバーで発言してみましょう
無事出力されています。
ここから、もしファイルだったらVirusTotalにスキャンリクエストを送るという挙動にしていきます。
ファイルだった場合%Alchemy.Messageのattachmentsに%Alchemy.Attachmentが入っているためそこの確認をして分岐したいと思います
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
こんな感じにすればいい感じにパターンマッチでうごくと思います。
良い感じです
ただ、これだと写真もスキャン対象になってしまいます。
なので正規表現を使って画像以外のスキャンをするように、しましょう。
ついでにトークンも追加
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を使うためのコードを書いていきます。
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するコード書いていきます。
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使ってるのでおしゃんてい
その後response_codeが1(成功)だったらetsにinsertみたいな感じです。
エラーだった時メッセージを変えるみたいなことしたかったのですが、12/24当日に書いてるので時間の都合上できませんでした
定期実行するためにquantumというライブラリが必要だったので適当に追加します
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を追加
defmodule VirusScanner.Scheduler do
use Quantum.Scheduler,
otp_app: :virus_scanner
end
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に全部書いていきます
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が入ってるサーバーに適当なファイルを上げてみてください。
いいいいかんじです
最後に
時間が無かったのでかなり適当に作りましたが割といい感じに動いたのでよかったです。
ソースは気が向いたらgithubに上げておきます。
では、クリスマスイブにプログラミンしてるぼくはケーキ食ってねます