More than 1 year has passed since last update.

DiscordAdvent Calendar 2017

@doge_boi

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

Last updated at 2022-02-18Posted at 2017-12-24

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

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

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

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

どんなの

ファイルがアップロードされたらそのファイルのスキャンリクエストを投げる
スキャンが終わったらそのファイルの結果をファイルが送られたチャンネルに送信する。

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

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

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

なぜ現実的ではないかと言うとAPI制限があるから
1に関しては 1分間に5個以上のファイルが投稿された場合待たないとファイルがダウンロードできない
2はファイル以上に投稿頻度が高いためスキャンが追い付かない

作ってく

言語はElixirで作ります。

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

~~今回は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が入ってるサーバーで発言してみましょう

無事出力されています。

ここから、もしファイルだったら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

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

良い感じです

ただ、これだと写真もスキャン対象になってしまいます。
なので正規表現を使って画像以外のスキャンをするように、しましょう。
ついでにトークンも追加

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使ってるのでおしゃんてい

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

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

定期実行するために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が入ってるサーバーに適当なファイルを上げてみてください。

いいいいかんじです

最後に

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

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

では、クリスマスイブにﾌﾟﾛｸﾞﾗﾐﾝしてるぼくはケーキ食ってねます

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up