/*
Technitium DNS Server
Copyright (C) 2021  Shreyas Zare (shreyas@technitium.com)

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.

*/

using DnsApplicationCommon;
using Microsoft.Data.Sqlite;
using Newtonsoft.Json;
using System;
using System.Collections.Generic;
using System.IO;
using System.Net;
using System.Threading;
using System.Threading.Tasks;
using TechnitiumLibrary.Net.Dns;

namespace QueryLogsSqlite
{
    public class App : IDnsLogger
    {
        #region variables

        bool _enableLogging;
        int _maxLogDays;
        string _connectionString;

        Timer _cleanupTimer;
        const int CLEAN_UP_TIMER_INITIAL_INTERVAL = 5 * 1000;
        const int CLEAN_UP_TIMER_PERIODIC_INTERVAL = 15 * 60 * 1000;

        #endregion

        #region IDisposable

        public void Dispose()
        {
            if (_cleanupTimer is not null)
            {
                _cleanupTimer.Dispose();
                _cleanupTimer = null;
            }
        }

        #endregion

        #region public

        public Task InitializeAsync(IDnsServer dnsServer, string config)
        {
            dynamic jsonConfig = JsonConvert.DeserializeObject(config);

            _enableLogging = jsonConfig.enableLogging.Value;
            _maxLogDays = Convert.ToInt32(jsonConfig.maxLogDays.Value);

            string sqliteDbPath = jsonConfig.sqliteDbPath.Value;
            string connectionString = jsonConfig.connectionString.Value;

            if (!Path.IsPathRooted(sqliteDbPath))
                sqliteDbPath = Path.Combine(dnsServer.ApplicationFolder, sqliteDbPath);

            _connectionString = connectionString.Replace("{sqliteDbPath}", sqliteDbPath);

            using (SqliteConnection connection = new SqliteConnection(_connectionString))
            {
                connection.Open();

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = @"
CREATE TABLE IF NOT EXISTS dns_logs
(
    dlid INTEGER PRIMARY KEY,
    timestamp DATETIME NOT NULL DEFAULT (DATETIME('now')),
    client_ip VARCHAR(39) NOT NULL,
    protocol TINYINT NOT NULL,
    response_type TINYINT NOT NULL,
    rcode TINYINT NOT NULL,
    qname VARCHAR(255),
    qtype SMALLINT,
    qclass SMALLINT,
    answer TEXT
);
";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_timestamp ON dns_logs (timestamp);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_client_ip ON dns_logs (client_ip);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_protocol ON dns_logs (protocol);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_response_type ON dns_logs (response_type);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_rcode ON dns_logs (rcode);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_qname ON dns_logs (qname);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_qtype ON dns_logs (qtype);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_qclass ON dns_logs (qclass);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_timestamp_client_ip ON dns_logs (timestamp, client_ip);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_timestamp_qname ON dns_logs (timestamp, qname);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_client_qname ON dns_logs (client_ip, qname);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_query ON dns_logs (qname, qtype);";
                    command.ExecuteNonQuery();
                }

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "CREATE INDEX IF NOT EXISTS index_all ON dns_logs (timestamp, client_ip, protocol, response_type, rcode, qname, qtype, qclass);";
                    command.ExecuteNonQuery();
                }
            }

            if (_maxLogDays < 1)
            {
                if (_cleanupTimer is not null)
                {
                    _cleanupTimer.Dispose();
                    _cleanupTimer = null;
                }
            }
            else
            {
                _cleanupTimer = new Timer(delegate (object state)
                {
                    try
                    {
                        using (SqliteConnection connection = new SqliteConnection(_connectionString))
                        {
                            connection.Open();

                            using (SqliteCommand command = connection.CreateCommand())
                            {
                                command.CommandText = "DELETE FROM dns_logs WHERE timestamp < @timestamp;";

                                command.Parameters.AddWithValue("@timestamp", DateTime.UtcNow.AddDays(_maxLogDays * -1));

                                command.ExecuteNonQuery();
                            }
                        }
                    }
                    catch (Exception ex)
                    {
                        dnsServer.WriteLog(ex);
                    }
                    finally
                    {
                        _cleanupTimer.Change(CLEAN_UP_TIMER_PERIODIC_INTERVAL, Timeout.Infinite);
                    }
                }, null, Timeout.Infinite, Timeout.Infinite);

                _cleanupTimer.Change(CLEAN_UP_TIMER_INITIAL_INTERVAL, Timeout.Infinite);
            }

            return Task.CompletedTask;
        }

        public Task InsertLogAsync(DnsDatagram request, IPEndPoint remoteEP, DnsTransportProtocol protocol, DnsDatagram response)
        {
            if (!_enableLogging)
                return Task.CompletedTask;

            DnsServerResponseType responseType;

            if (response.Tag == null)
                responseType = DnsServerResponseType.Recursive;
            else
                responseType = (DnsServerResponseType)response.Tag;

            DnsQuestionRecord query;

            if (request.Question.Count > 0)
                query = request.Question[0];
            else
                query = null;

            string answer = null;

            if (response is null)
            {
                answer = null;
            }
            else if (response.Answer.Count == 0)
            {
                answer = "";
            }
            else if ((response.Answer.Count > 2) && response.IsZoneTransfer)
            {
                answer = "[ZONE TRANSFER]";
            }
            else
            {
                for (int i = 0; i < response.Answer.Count; i++)
                {
                    if (answer is null)
                        answer = response.Answer[i].RDATA.ToString();
                    else
                        answer += ", " + response.Answer[i].RDATA.ToString();
                }
            }

            using (SqliteConnection connection = new SqliteConnection(_connectionString))
            {
                connection.Open();

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "INSERT INTO dns_logs (client_ip, protocol, response_type, rcode, qname, qtype, qclass, answer) VALUES (@client_ip, @protocol, @response_type, @rcode, @qname, @qtype, @qclass, @answer);";

                    command.Parameters.AddWithValue("@client_ip", remoteEP.Address.ToString());
                    command.Parameters.AddWithValue("@protocol", (byte)protocol);
                    command.Parameters.AddWithValue("@response_type", (byte)responseType);
                    command.Parameters.AddWithValue("@rcode", (byte)response.RCODE);

                    if (query is null)
                    {
                        command.Parameters.AddWithValue("@qname", DBNull.Value);
                        command.Parameters.AddWithValue("@qtype", DBNull.Value);
                        command.Parameters.AddWithValue("@qclass", DBNull.Value);
                    }
                    else
                    {
                        command.Parameters.AddWithValue("@qname", query.Name);
                        command.Parameters.AddWithValue("@qtype", (ushort)query.Type);
                        command.Parameters.AddWithValue("@qclass", (ushort)query.Class);
                    }

                    if (answer is null)
                        command.Parameters.AddWithValue("@answer", DBNull.Value);
                    else
                        command.Parameters.AddWithValue("@answer", answer);

                    command.ExecuteNonQuery();
                }
            }

            return Task.CompletedTask;
        }

        public Task<DnsLogPage> QueryLogsAsync(long pageNumber, int entriesPerPage, DateTime? start, DateTime? end, IPAddress clientIpAddress, DnsTransportProtocol? protocol, DnsServerResponseType? responseType, DnsResponseCode? rcode, string qname, DnsResourceRecordType? qtype, DnsClass? qclass)
        {
            string whereClause = string.Empty;

            if (start is not null)
                whereClause += "timestamp >= @start AND ";

            if (end is not null)
                whereClause += "timestamp <= @end AND ";

            if (clientIpAddress is not null)
                whereClause += "client_ip = @client_ip AND ";

            if (protocol is not null)
                whereClause += "protocol = @protocol AND ";

            if (responseType is not null)
                whereClause += "response_type = @response_type AND ";

            if (rcode is not null)
                whereClause += "rcode = @rcode AND ";

            if (qname is not null)
                whereClause += "qname = @qname AND ";

            if (qtype is not null)
                whereClause += "qtype = @qtype AND ";

            if (qclass is not null)
                whereClause += "qclass = @qclass AND ";

            if (!string.IsNullOrEmpty(whereClause))
                whereClause = whereClause.Substring(0, whereClause.Length - 5);

            using (SqliteConnection connection = new SqliteConnection(_connectionString))
            {
                connection.Open();

                //find total entries
                long totalEntries;

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = "SELECT Count(*) FROM dns_logs" + (string.IsNullOrEmpty(whereClause) ? ";" : " WHERE " + whereClause + ";");

                    if (start is not null)
                        command.Parameters.AddWithValue("@start", start);

                    if (end is not null)
                        command.Parameters.AddWithValue("@end", end);

                    if (clientIpAddress is not null)
                        command.Parameters.AddWithValue("@client_ip", clientIpAddress.ToString());

                    if (protocol is not null)
                        command.Parameters.AddWithValue("@protocol", (byte)protocol);

                    if (responseType is not null)
                        command.Parameters.AddWithValue("@response_type", (byte)responseType);

                    if (rcode is not null)
                        command.Parameters.AddWithValue("@rcode", (byte)rcode);

                    if (qname is not null)
                        command.Parameters.AddWithValue("@qname", qname);

                    if (qtype is not null)
                        command.Parameters.AddWithValue("@qtype", (ushort)qtype);

                    if (qclass is not null)
                        command.Parameters.AddWithValue("@qclass", (ushort)qclass);

                    totalEntries = (long)command.ExecuteScalar();
                }

                long totalPages = (totalEntries / entriesPerPage) + (totalEntries % entriesPerPage > 0 ? 1 : 0);

                if (pageNumber > totalPages)
                    pageNumber = totalPages;

                long endRowNum = pageNumber * entriesPerPage;
                long startRowNum = endRowNum - entriesPerPage;

                List<DnsLogEntry> entries = new List<DnsLogEntry>(entriesPerPage);

                using (SqliteCommand command = connection.CreateCommand())
                {
                    command.CommandText = @"
SELECT * FROM (
    SELECT
        ROW_NUMBER() OVER ( 
            ORDER BY timestamp
        ) row_num,
        timestamp,
        client_ip,
        protocol,
        response_type,
        rcode,
        qname,
        qtype,
        qclass,
        answer
    FROM
        dns_logs
" + (string.IsNullOrEmpty(whereClause) ? "" : "WHERE " + whereClause) + @"
) t
WHERE 
    row_num > @start_row_num AND row_num <= @end_row_num
";

                    command.Parameters.AddWithValue("@start_row_num", startRowNum);
                    command.Parameters.AddWithValue("@end_row_num", endRowNum);

                    if (start is not null)
                        command.Parameters.AddWithValue("@start", start);

                    if (end is not null)
                        command.Parameters.AddWithValue("@end", end);

                    if (clientIpAddress is not null)
                        command.Parameters.AddWithValue("@client_ip", clientIpAddress.ToString());

                    if (protocol is not null)
                        command.Parameters.AddWithValue("@protocol", (byte)protocol);

                    if (responseType is not null)
                        command.Parameters.AddWithValue("@response_type", (byte)responseType);

                    if (rcode is not null)
                        command.Parameters.AddWithValue("@rcode", (byte)rcode);

                    if (qname is not null)
                        command.Parameters.AddWithValue("@qname", qname);

                    if (qtype is not null)
                        command.Parameters.AddWithValue("@qtype", (ushort)qtype);

                    if (qclass is not null)
                        command.Parameters.AddWithValue("@qclass", (ushort)qclass);

                    using (SqliteDataReader reader = command.ExecuteReader())
                    {
                        while (reader.Read())
                        {
                            DnsQuestionRecord question;

                            if (reader.IsDBNull(6))
                                question = null;
                            else
                                question = new DnsQuestionRecord(reader.GetString(6), (DnsResourceRecordType)reader.GetInt16(7), (DnsClass)reader.GetInt16(8));

                            string answer;

                            if (reader.IsDBNull(9))
                                answer = null;
                            else
                                answer = reader.GetString(9);

                            entries.Add(new DnsLogEntry(reader.GetInt64(0), reader.GetDateTime(1), IPAddress.Parse(reader.GetString(2)), (DnsTransportProtocol)reader.GetByte(3), (DnsServerResponseType)reader.GetByte(4), (DnsResponseCode)reader.GetByte(5), question, answer));
                        }
                    }
                }

                return Task.FromResult(new DnsLogPage(pageNumber, totalPages, totalEntries, entries));
            }
        }

        #endregion

        #region properties

        public string Description
        { get { return "Logs all incoming DNS requests and their responses in a Sqlite database that can be queried from the DNS Server web console."; } }

        #endregion
    }
}