/*
    Copyright 2016-2019 Gleb Popov <6yearold@gmail.com>
    Copyright (c) 2023 Serenity Cyber Security, LLC <license@futurecrew.ru>
                Author: Alexey Yushkin <x636808@gmail.com>

    Redistribution and use in source and binary forms, with or without modification,
    are permitted provided that the following conditions are met:
    1. Redistributions of source code must retain the above copyright notice,
    this list of conditions and the following disclaimer.
    2. Redistributions in binary form must reproduce the above copyright notice,
    this list of conditions and the following disclaimer in the documentation
    and/or other materials provided with the distribution.
    3. Neither the name of the copyright holder nor the names of its contributors may
    be used to endorse or promote products derived from this software without specific
    prior written permission.

    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
    ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
    WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
    IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
    INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
    BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
    LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
    OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
    ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <ranges>

#include "bsdisks.h"
#include "adaptors.h"
#include "objectmanager.h"
#include "block.h"
#include "blockfilesystem.h"
#include "drive.h"
#include "filesystemprober.h"
#include "camcontrolprober.h"
#include "zfsprober.h"

// We have to run this check whenever a probe is finished, because the object
// being probed may be deleted meanwhile.
#define GET_BLOCK_OR_RETURN_IF_GONE(name) Block* b = m_blockObjects.value(name, nullptr); \
    if(!b) return;

namespace views = std::views;

static QString& normalizeForDBusPath(QString& s)
{
    // dbus paths cannot contain '+', which may come from logical partitions, i.e ada0s2+12345
    return s.replace(QChar('+'), QChar('_'));
}

static ObjectManager* TheObjectManager;

ObjectManager::ObjectManager()
{
    TheObjectManager = this;
}

ObjectManager& ObjectManager::Instance()
{
    return *TheObjectManager;
}

DBUSManagerStruct ObjectManager::GetManagedObjects()
{
    DBUSManagerStruct ret;

    QList<std::pair<QObject*, QDBusObjectPath>> objects;

    std::transform(m_driveObjects.begin(), m_driveObjects.end(), std::back_inserter(objects),
        [](Drive* d) { return std::make_pair(static_cast<QObject*>(d), d->dbusPath); }
    );

    std::transform(m_blockObjects.begin(), m_blockObjects.end(), std::back_inserter(objects),
        [](Block* b) { return std::make_pair(static_cast<QObject*>(b), b->dbusPath); }
    );

    for(const auto& objectPair : objects)
    {
        QObject* o = objectPair.first;
        auto dbusPath = objectPair.second;

        QVariantMapMap interfaces;

        for(QObject* child : o->children())
        {
            QDBusAbstractAdaptor* adaptor = qobject_cast<QDBusAbstractAdaptor*>(child);

            if(!adaptor)
                continue;

            const QString& iface = adaptor->metaObject()->classInfo(adaptor->metaObject()->indexOfClassInfo("D-Bus Interface")).value();

            QVariantMap properties;
            for (int i = adaptor->metaObject()->propertyOffset();
                    i < adaptor->metaObject()->propertyCount(); ++i) {
                auto propertyName = adaptor->metaObject()->property(i).name();
                properties.insert(QString::fromLatin1(propertyName), adaptor->property(propertyName));
            }

            interfaces[iface] = properties;
        }

        ret[dbusPath] = interfaces;
    }

    return ret;
}


void ObjectManager::initialProbe()
{
    auto* geomProber = new GeomProber;

    QObject::connect(geomProber, &GeomProber::gotDisk, this,
        [this](DiskInfo di)
        {
            addDrive(di.devName());

            // assume that any object from `geom disk list` has a block device in /dev
            addBlock(di.devName());
        });
    QObject::connect(geomProber, &GeomProber::gotPartTable, this,
        [this](PartTableInfo pi)
        {
            // memory disks are not reported by `geom disk list`, so we
            // have to add them here
            if(pi.geomName().startsWith("md"))
                addBlock(pi.geomName());

            for(const auto& devName : pi.partitions().keys())
                addBlock(devName);
        });

    QThreadPool::globalInstance()->start(geomProber);

    if (BsdisksConfig::get().HandleZFSFilesystems)
    {
        auto* zfsProber = new ZFSProber();
        QObject::connect(zfsProber, &ZFSProber::gotDataset, this, &ObjectManager::addZFSDataset);

        QThreadPool::globalInstance()->start(zfsProber);
    }

    initialProbeDone = true;
}

void ObjectManager::CheckAuth(const QString& actionId, QObject* subject, AuthQueue::Continuation continuation)
{
    m_authQueue.Enqueue(actionId, subject, continuation);
}

bool ObjectManager::BlockObjectValid(Block* block) const
{
    return !m_blockObjects.key(block).isEmpty();
}

bool ObjectManager::DriveObjectValid(Drive* drive) const
{
    return !m_driveObjects.key(drive).isEmpty();
}

void ObjectManager::filesystemAdded(Block* b, QString fs, const ZFSInfo* zfsInfo)
{
    Q_ASSERT(b->bFilesystem == nullptr);

    auto bfs = new BlockFilesystem(b);
    bfs->filesystem = fs;

    b->bFilesystem = bfs;
    b->idUsage = QStringLiteral("filesystem");

    if(fs == QStringLiteral("msdosfs"))
    {
        b->idType = QStringLiteral("vfat");
    }
    else
    {
        b->idType = fs;
    }

    if(fs == QStringLiteral("zfs"))
    {
        bfs->zfsDataset = zfsInfo->dataSet();
        bfs->zfsMountpoint = zfsInfo->mountPoint();
        if(zfsInfo->isMounted())
            bfs->mountPoints << qMakePair(zfsInfo->mountPoint().toLocal8Bit(), 0);

        b->hasNoDrive = true;
    }
    else {} // fill non-ZFS mount points in registerBlock, when all probes are done

    if (b->registered)
    {
        // Filesystem appeared after registration due to call to Rescan()
        QList<std::pair<QString, QDBusAbstractAdaptor*>> interfaces;
        interfaces << std::make_pair(QStringLiteral("org.freedesktop.UDisks2.Filesystem"), new FilesystemAdaptor(b));
        addInterfaces(b->dbusPath, interfaces);
    }
}

void ObjectManager::addBlock(QString dev)
{
    if(auto* drive = m_driveObjects.value(dev))
    {
        if(drive->vendor() == "VBOX CD-ROM" && m_blockObjects.contains(dev))
            return;
    }

    // don't do anything with stuff like diskid/blabla
    if (dev.contains('/'))
        return;

    Q_ASSERT(!m_blockObjects.contains(dev));

    qDebug() << "Created block " << dev;
    auto* b = new Block;
    m_blockObjects.insert(dev, b);

    b->name = dev;
    b->dbusPath = QDBusObjectPath(UDisksBlockDevices + normalizeForDBusPath(dev));

    startFilesystemProbe(b);
    startGeomProbe(b);
}

void ObjectManager::updateBlock(QString dev)
{
    // don't do anything with stuff like diskid/blabla
    if (dev.contains('/'))
        return;

    auto* b = m_blockObjects.value(dev, nullptr);
    // the device gone somewhere between receiving a notification and this call
    if(!b)
        return;

    if(b->probesDone.all())
    {
        b->probesDone.reset();
        b->probesDone.set(FILESYSTEM_PROBE);
        startGeomProbe(b);
    }
    else
        b->needsAnotherProbe = true;
}


void ObjectManager::removeBlock(QString dev)
{
    if(auto* drive = m_driveObjects.value(dev))
    {
        if(drive->vendor() == "VBOX CD-ROM")
            return;
    }

    // happens with partitions blocks - we delete them along with table block,
    // but the devd event may arrive after that
    if(!m_blockObjects.contains(dev))
        return;

    auto* b = m_blockObjects.take(dev);
    auto* d = m_driveObjects.value(b->driveName());
    qDebug() << "Unregistering " + b->dbusPath.path();
    QStringList ifaces;

    if(b->bFilesystem)
        ifaces << QStringLiteral("org.freedesktop.UDisks2.Filesystem");
    if(b->bPartition)
        ifaces << QStringLiteral("org.freedesktop.UDisks2.Partition");
    if(b->bPartTable)
    {
        ifaces << QStringLiteral("org.freedesktop.UDisks2.PartitionTable");
        for(const auto& pb : b->bPartTable->partitionBlockNames)
            removeBlock(pb);
    }
    ifaces << QStringLiteral("org.freedesktop.UDisks2.Block");
    removeInterfaces(b->dbusPath, ifaces);
    QDBusConnection::systemBus().unregisterObject(b->dbusPath.path());

    if(!b->hasNoDrive)
        // delete drive only when we are deleting table block or any other non-partition block
        if(!b->bPartition)
            removeDrive(d->geomName);

    delete b;
}

void ObjectManager::addDrive(QString dev)
{
    Q_ASSERT(!m_driveObjects.contains(dev));

    qDebug() << "Created drive " << dev;
    auto* d = new Drive;
    m_driveObjects.insert(dev, d);

    QString devPath = UDisksDrives + dev;
    d->dbusPath = QDBusObjectPath(devPath);
    d->geomName = dev;

    auto geomProber = new GeomProber(d->geomName);
    QObject::connect(geomProber, &GeomProber::gotDisk, this,
                     [this, d](DiskInfo di)
    {
        d->size = di.mediaSize();
        d->description = di.descr();
        d->geomIdentifier = di.ident();

        d->geomProbeDone = true;
        if(d->camcontrolProbeDone)
            registerDrive(d);
    });

    auto camControlProber = new CamControlProber(d->geomName);
    QObject::connect(camControlProber, &CamControlProber::finished, this,
                     [this, d](const QString& serial, bool isRemovable, int transport, bool isAtaSata)
    {
        d->camIdentifier = serial;
        d->isRemovable = isRemovable;
        d->ataSata = isAtaSata ? QStringLiteral("sata") : QStringLiteral("ata");
        d->transport = static_cast<camtransport>(transport);

        d->camcontrolProbeDone = true;
        if(d->geomProbeDone)
            registerDrive(d);
    });

    QThreadPool::globalInstance()->start(geomProber);
    QThreadPool::globalInstance()->start(camControlProber);
}

void ObjectManager::removeDrive(QString dev)
{
    Q_ASSERT(m_driveObjects.contains(dev));

    auto* d = m_driveObjects.take(dev);
    qDebug() << "Unregistering " + d->dbusPath.path();
    removeInterfaces(d->dbusPath, { QStringLiteral("org.freedesktop.UDisks2.Drive") });
    QDBusConnection::systemBus().unregisterObject(d->dbusPath.path());
    delete d;
}

void ObjectManager::postponeRegistration(QString blockName)
{
    m_postponedRegistrations.insert(blockName);
}

void ObjectManager::tryRegisterPostponed()
{
    int postponedSize;
    do
    {
        postponedSize = m_postponedRegistrations.size();

        for(auto it = m_postponedRegistrations.begin(); it != m_postponedRegistrations.end(); it++)
        {
            if(registerBlock(m_blockObjects[*it], false))
            {
                qDebug() << "Pop " << *it << " from m_postponedRegistrations";
                m_postponedRegistrations.erase(it);
                break;
            }
        }
    } while(postponedSize > m_postponedRegistrations.size());
}

bool ObjectManager::registerBlock(Block* b, bool tryPostponed)
{
    // We've got MEDIACHANGE event while probes were running.
    // Re-start GEOM probe and postpone registration.
    if(b->needsAnotherProbe)
    {
        qDebug() << b->name << " changed, restarting probes ";
        b->needsAnotherProbe = false;
        b->probesDone.reset(GEOM_PROBE);
        startGeomProbe(b);
        return false;
    }

    if(tryPostponed)
        tryRegisterPostponed();

    if(b->bPartition)
    {
        Block* table = m_blockObjects.value(b->bPartition->partBlockName, nullptr);
        if(!table || !table->registered)
        {
            qDebug() << b->name << " waits for partition table " << (table ? table->name : QString());
            postponeRegistration(b->name);
            return false;
        }
        if(!b->bPartition->partTableBlock)
            b->bPartition->partTableBlock = table;
    }

    if(!b->hasNoDrive)
        if(Drive* d = m_driveObjects.value(b->driveName(), nullptr))
        {
            if(!d || !(d->camcontrolProbeDone && d->geomProbeDone))
            {
                qDebug() << b->name << " waits for drive " << b->driveName();
                postponeRegistration(b->name);
                return false;
            }
        }

    if(!b->registered)
    {
        if(b->bFilesystem)
            b->fillMountPoints();

        b->registered = true;
        b->parentDrive = m_driveObjects.value(b->driveName(), nullptr);

        QString devPath = b->dbusPath.path();

        qDebug() << "Registering " + devPath;
        QDBusConnection::systemBus().registerObject(devPath, b);

        QList<std::pair<QString, QDBusAbstractAdaptor*>> interfaces;
        interfaces << std::make_pair(QStringLiteral("org.freedesktop.UDisks2.Block"), new BlockAdaptor(b));

        if(b->bFilesystem)
            interfaces << std::make_pair(QStringLiteral("org.freedesktop.UDisks2.Filesystem"), new FilesystemAdaptor(b));

        if(b->bPartTable)
            interfaces << std::make_pair(QStringLiteral("org.freedesktop.UDisks2.PartitionTable"), new PartitionTableAdaptor(b));

        if(b->bPartition)
            interfaces << std::make_pair(QStringLiteral("org.freedesktop.UDisks2.Partition"), new PartitionAdaptor(b));

        addInterfaces(b->dbusPath, interfaces);
    }
    else {} // we were updating things after updateBlock()

    if(tryPostponed)
        tryRegisterPostponed();

    return true;
}


void ObjectManager::registerDrive(Drive* d)
{
    const QString dbusPath = d->dbusPath.path();

    qDebug() << "Registering " + dbusPath;
    QDBusConnection::systemBus().registerObject(dbusPath, d);

    addInterfaces(d->dbusPath, {std::make_pair("org.freedesktop.UDisks2.Drive", new DriveAdaptor(d))});

    tryRegisterPostponed();
}

void ObjectManager::addZFSDataset(const ZFSInfo& zfsInfo)
{
    QString normalizedDatasetName = zfsInfo.dataSet()
                                           .replace('/', '_')  // turn pool/foo into pool_foo
                                           .replace('.', '_'); // D-Bus dislikes dots in paths
    Q_ASSERT(!m_blockObjects.contains(normalizedDatasetName));

    auto* b = new Block;
    m_blockObjects.insert(normalizedDatasetName, b);

    QString devPath = UDisksBlockDevices + normalizedDatasetName;
    b->name = normalizedDatasetName;
    b->dbusPath = QDBusObjectPath(devPath);

    filesystemAdded(b, QStringLiteral("zfs"), &zfsInfo);

    registerBlock(b);
}

void ObjectManager::mountFS(QString mountDevice, QString mountPoint)
{
    auto blocks = m_blockObjects.values()
        | views::filter([](const auto *block) { return block->bFilesystem; });
    for(const auto* block: blocks)
        block->bFilesystem->tryAddMountPoint(mountDevice, mountPoint);
}

void ObjectManager::unmountFS(QString mountPoint)
{
    for(auto* b : m_blockObjects.values())
    {
        auto* bfs = b->bFilesystem;

        if(!bfs)
            continue;

        bfs->tryRemoveMountPoint(mountPoint);
    }
}

void ObjectManager::startFilesystemProbe(Block* b)
{
    auto prober = new FilesystemProber(b->device());
    const QString devName = b->name;

    QObject::connect(prober, &FilesystemProber::finished, this,
                     [this, devName](QString fs)
    {
        GET_BLOCK_OR_RETURN_IF_GONE(devName);

        qDebug() << "Finished FS probe on " << b->name;
        if(!fs.isEmpty())
            filesystemAdded(b, fs);

        b->probesDone.set(FILESYSTEM_PROBE);
        if(b->probesDone.all())
            registerBlock(b);
    });

    QThreadPool::globalInstance()->start(prober);
}


void ObjectManager::startGeomProbe(Block* b)
{
    auto prober = new GeomProber(b->name);

    if(b->name.startsWith("md"))
        b->hasNoDrive = true;

    const QString devName = b->name;

    QObject::connect(prober, &GeomProber::gotLabels, this,
                     [this, devName](QStringList labels)
    {
        GET_BLOCK_OR_RETURN_IF_GONE(devName);

        b->labels = labels;
    });
    QObject::connect(prober, &GeomProber::gotDisk, this,
                     [this, devName](DiskInfo di)
    {
        GET_BLOCK_OR_RETURN_IF_GONE(devName);

        b->size = di.mediaSize();
        b->description = di.descr();
        b->identifier = di.ident();

        if(!b->hasNoDrive && !m_driveObjects.contains(devName))
        {
            addDrive(devName);
        }
    });
    QObject::connect(prober, &GeomProber::gotPartTable, this,
                     [this, devName](PartTableInfo pi)
    {
        GET_BLOCK_OR_RETURN_IF_GONE(devName);

        // do not create BlockPartTable when called from updateBlock()
        if(!b->bPartTable)
            b->bPartTable = new BlockPartTable(b);

        b->bPartTable->setTableType(pi.scheme());

        for(auto it = pi.partitions().constBegin(); it != pi.partitions().cend(); it++)
            b->bPartTable->partitionBlockNames << it.key();
    });
    QObject::connect(prober, &GeomProber::gotPart, this,
                     [this, devName](QString tableBlockName, Part partInfo)
    {
        GET_BLOCK_OR_RETURN_IF_GONE(devName);

        addPartition(b, tableBlockName, partInfo);
    });
    QObject::connect(prober, &GeomProber::finished, this,
                     [this, devName]()
    {
        GET_BLOCK_OR_RETURN_IF_GONE(devName);

        b->probesDone.set(GEOM_PROBE);
        qDebug() << "Finished GEOM probe on " << b->name;
        if(b->probesDone.all())
            registerBlock(b);
    });

    QThreadPool::globalInstance()->start(prober);
}

void ObjectManager::addPartition(Block* b, const QString& tableBlockName, const Part& partInfo)
{
    auto* bp = new BlockPartition(b);

    bp->partBlockName = tableBlockName;
    bp->partitionType = partInfo.type();
    bp->number = partInfo.index();
    bp->offset = partInfo.offset();
    bp->size = partInfo.length();
    bp->uuid = partInfo.uuid();

    b->bPartition = bp;
}


void ObjectManager::addInterfaces(QDBusObjectPath path, QList<std::pair<QString, QDBusAbstractAdaptor*>> newInterfaces)
{
    QVariantMapMap interfaces;

    for (const auto& pair : newInterfaces)
    {
        const QString& iface = pair.first;
        QDBusAbstractAdaptor* adaptor = pair.second;

        QVariantMap properties;
        for (int i = adaptor->metaObject()->propertyOffset();
                i < adaptor->metaObject()->propertyCount(); ++i) {
            auto propertyName = adaptor->metaObject()->property(i).name();
            properties.insert(QString::fromLatin1(propertyName), adaptor->property(propertyName));
        }

        interfaces[iface] = properties;
    }

    emit InterfacesAdded(path, interfaces);
}


void ObjectManager::removeInterfaces(QDBusObjectPath path, QStringList ifaces)
{
    emit InterfacesRemoved(path, ifaces);
}
