_loopback.cpp

/**
 * JUPITER Sound Booster — WASAPI Process Loopback Extension
 *
 * Captures ALL system audio EXCEPT our own process output.
 * Uses ActivateAudioInterfaceAsync with PROCESS_LOOPBACK_MODE_EXCLUDE_TARGET_PROCESS_TREE.
 * Requires Windows 10 Build 20438+ (version 2004/20H1).
 *
 * Based on Microsoft's Application Loopback Audio Capture sample
 * and ProcTap (MIT license, github.com/m96-chan/ProcTap).
 */

#define PY_SSIZE_T_CLEAN
#include <Python.h>
#include <windows.h>
#include <mmdeviceapi.h>
#include <audioclient.h>
#include <audiopolicy.h>
#include <combaseapi.h>
#include <objidl.h>
#include <wrl/client.h>
#include <vector>

using Microsoft::WRL::ComPtr;

/* ── Constants & structs (may be missing from older SDK headers) ── */

#ifndef VIRTUAL_AUDIO_DEVICE_PROCESS_LOOPBACK
#define VIRTUAL_AUDIO_DEVICE_PROCESS_LOOPBACK L"VAD\\Process_Loopback"
#endif

enum AUDIOCLIENT_ACTIVATION_TYPE {
    AUDIOCLIENT_ACTIVATION_TYPE_DEFAULT = 0,
    AUDIOCLIENT_ACTIVATION_TYPE_PROCESS_LOOPBACK = 1
};

enum PROCESS_LOOPBACK_MODE {
    PROCESS_LOOPBACK_MODE_INCLUDE_TARGET_PROCESS_TREE = 0,
    PROCESS_LOOPBACK_MODE_EXCLUDE_TARGET_PROCESS_TREE = 1
};

struct AUDIOCLIENT_PROCESS_LOOPBACK_PARAMS {
    DWORD TargetProcessId;
    PROCESS_LOOPBACK_MODE ProcessLoopbackMode;
};

struct AUDIOCLIENT_ACTIVATION_PARAMS {
    AUDIOCLIENT_ACTIVATION_TYPE ActivationType;
    union {
        AUDIOCLIENT_PROCESS_LOOPBACK_PARAMS ProcessLoopbackParams;
    };
};

typedef HRESULT(WINAPI* PFN_ActivateAudioInterfaceAsync)(
    LPCWSTR deviceInterfacePath,
    REFIID riid,
    PROPVARIANT* activationParams,
    IActivateAudioInterfaceCompletionHandler* completionHandler,
    IActivateAudioInterfaceAsyncOperation** activationOperation
);

/* ── Async completion handler (COM callback) ── */

class ActivationHandler : public IActivateAudioInterfaceCompletionHandler,
                          public IAgileObject {
    LONG m_ref;
    HANDLE m_event;
    HRESULT m_hr;
    ComPtr<IUnknown> m_result;
public:
    ActivationHandler() : m_ref(1), m_hr(E_FAIL) {
        m_event = CreateEvent(nullptr, TRUE, FALSE, nullptr);
    }
    ~ActivationHandler() { if (m_event) CloseHandle(m_event); }

    STDMETHODIMP QueryInterface(REFIID riid, void** ppv) override {
        if (riid == __uuidof(IUnknown) ||
            riid == __uuidof(IActivateAudioInterfaceCompletionHandler)) {
            *ppv = static_cast<IActivateAudioInterfaceCompletionHandler*>(this);
            AddRef(); return S_OK;
        }
        if (riid == __uuidof(IAgileObject)) {
            *ppv = static_cast<IAgileObject*>(this);
            AddRef(); return S_OK;
        }
        *ppv = nullptr; return E_NOINTERFACE;
    }
    STDMETHODIMP_(ULONG) AddRef()  override { return InterlockedIncrement(&m_ref); }
    STDMETHODIMP_(ULONG) Release() override {
        LONG c = InterlockedDecrement(&m_ref);
        if (c == 0) delete this;
        return c;
    }
    STDMETHODIMP ActivateCompleted(IActivateAudioInterfaceAsyncOperation* op) override {
        HRESULT hrResult = E_FAIL;
        ComPtr<IUnknown> pUnk;
        op->GetActivateResult(&hrResult, &pUnk);
        m_hr = hrResult;
        m_result = pUnk;
        SetEvent(m_event);
        return S_OK;
    }
    HRESULT Wait(DWORD ms = 10000) {
        if (WaitForSingleObject(m_event, ms) != WAIT_OBJECT_0) return E_FAIL;
        return m_hr;
    }
    ComPtr<IUnknown> GetResult() { return m_result; }
};

/* ── Main capture class ── */

class SystemLoopback {
    ComPtr<IAudioClient>        m_client;
    ComPtr<IAudioCaptureClient> m_capture;
    WAVEFORMATEX*               m_fmt = nullptr;
    bool                        m_running = false;
    bool                        m_processSpecific = false;
    std::vector<BYTE>           m_buf;
    CRITICAL_SECTION            m_lock;

public:
    SystemLoopback()  { InitializeCriticalSection(&m_lock); }
    ~SystemLoopback() { Cleanup(); DeleteCriticalSection(&m_lock); }

    /* Try process-specific loopback; fall back to system-wide */
    HRESULT Init() {
        HRESULT hr = CoInitializeEx(nullptr, COINIT_APARTMENTTHREADED);
        if (hr == RPC_E_CHANGED_MODE) hr = S_OK;  /* already init'd */
        if (FAILED(hr)) return hr;

        hr = InitExcludeProcess();
        if (SUCCEEDED(hr)) return hr;

        /* Fallback: system-wide loopback (will include our output — echo cancel needed) */
        return InitSystemWide();
    }

    HRESULT Start() {
        if (m_running) return S_OK;
        HRESULT hr = m_client->Start();
        if (SUCCEEDED(hr)) m_running = true;
        return hr;
    }

    void Stop() {
        if (!m_running) return;
        m_running = false;
        if (m_client) m_client->Stop();
    }

    /* Drain WASAPI packets into internal buffer, return all buffered data */
    PyObject* Read() {
        if (!m_running || !m_capture) Py_RETURN_NONE;

        /* Drain all available packets */
        UINT32 pktLen = 0;
        while (SUCCEEDED(m_capture->GetNextPacketSize(&pktLen)) && pktLen > 0) {
            BYTE* pData = nullptr;
            UINT32 frames = 0;
            DWORD flags = 0;
            if (FAILED(m_capture->GetBuffer(&pData, &frames, &flags, nullptr, nullptr)))
                break;
            UINT32 bytes = frames * m_fmt->nBlockAlign;
            EnterCriticalSection(&m_lock);
            size_t old = m_buf.size();
            m_buf.resize(old + bytes);
            if (flags & AUDCLNT_BUFFERFLAGS_SILENT)
                memset(m_buf.data() + old, 0, bytes);
            else
                memcpy(m_buf.data() + old, pData, bytes);
            LeaveCriticalSection(&m_lock);
            m_capture->ReleaseBuffer(frames);
        }

        EnterCriticalSection(&m_lock);
        if (m_buf.empty()) {
            LeaveCriticalSection(&m_lock);
            return PyBytes_FromStringAndSize("", 0);
        }
        PyObject* result = PyBytes_FromStringAndSize(
            (const char*)m_buf.data(), (Py_ssize_t)m_buf.size());
        m_buf.clear();
        LeaveCriticalSection(&m_lock);
        return result;
    }

    PyObject* GetFormat() {
        if (!m_fmt) Py_RETURN_NONE;
        return Py_BuildValue("{s:i,s:i,s:i,s:i}",
            "channels",       (int)m_fmt->nChannels,
            "sample_rate",    (int)m_fmt->nSamplesPerSec,
            "bits_per_sample",(int)m_fmt->wBitsPerSample,
            "block_align",    (int)m_fmt->nBlockAlign);
    }

    bool IsProcessSpecific() { return m_processSpecific; }

    void Cleanup() {
        Stop();
        m_capture.Reset();
        m_client.Reset();
        if (m_fmt) { CoTaskMemFree(m_fmt); m_fmt = nullptr; }
        EnterCriticalSection(&m_lock);
        m_buf.clear();
        LeaveCriticalSection(&m_lock);
    }

private:
    HRESULT InitExcludeProcess() {
        /* Load ActivateAudioInterfaceAsync dynamically */
        HMODULE hMod = LoadLibraryW(L"mmdevapi.dll");
        if (!hMod) return E_FAIL;

        auto pfn = (PFN_ActivateAudioInterfaceAsync)
            GetProcAddress(hMod, "ActivateAudioInterfaceAsync");
        if (!pfn) { FreeLibrary(hMod); return E_FAIL; }

        /* Activation params: EXCLUDE our own process tree */
        AUDIOCLIENT_ACTIVATION_PARAMS params = {};
        params.ActivationType = AUDIOCLIENT_ACTIVATION_TYPE_PROCESS_LOOPBACK;
        params.ProcessLoopbackParams.ProcessLoopbackMode =
            PROCESS_LOOPBACK_MODE_EXCLUDE_TARGET_PROCESS_TREE;
        params.ProcessLoopbackParams.TargetProcessId = GetCurrentProcessId();

        PROPVARIANT pv = {};
        pv.vt = VT_BLOB;
        pv.blob.cbSize = sizeof(params);
        pv.blob.pBlobData = (BYTE*)&params;

        auto* handler = new ActivationHandler();
        ComPtr<IActivateAudioInterfaceAsyncOperation> asyncOp;
        HRESULT hr = pfn(VIRTUAL_AUDIO_DEVICE_PROCESS_LOOPBACK,
                         __uuidof(IAudioClient), &pv, handler, &asyncOp);
        FreeLibrary(hMod);
        if (FAILED(hr)) { handler->Release(); return hr; }

        hr = handler->Wait(10000);
        if (FAILED(hr)) { handler->Release(); return hr; }

        ComPtr<IUnknown> pUnk = handler->GetResult();
        handler->Release();
        if (!pUnk) return E_FAIL;

        hr = pUnk.As(&m_client);
        if (FAILED(hr)) return hr;

        m_processSpecific = true;

        /* Process loopback: GetMixFormat may return E_NOTIMPL.
           Use hardcoded 48kHz/float32/stereo (standard). */
        m_fmt = (WAVEFORMATEX*)CoTaskMemAlloc(sizeof(WAVEFORMATEX));
        if (!m_fmt) return E_OUTOFMEMORY;
        m_fmt->wFormatTag      = WAVE_FORMAT_IEEE_FLOAT;
        m_fmt->nChannels       = 2;
        m_fmt->nSamplesPerSec  = 48000;
        m_fmt->wBitsPerSample  = 32;
        m_fmt->nBlockAlign     = m_fmt->nChannels * m_fmt->wBitsPerSample / 8;
        m_fmt->nAvgBytesPerSec = m_fmt->nSamplesPerSec * m_fmt->nBlockAlign;
        m_fmt->cbSize          = 0;

        hr = m_client->Initialize(AUDCLNT_SHAREMODE_SHARED,
                                  AUDCLNT_STREAMFLAGS_LOOPBACK,
                                  10000000, 0, m_fmt, nullptr);
        if (FAILED(hr)) {
            /* Fallback: 44.1kHz int16 */
            m_fmt->wFormatTag      = WAVE_FORMAT_PCM;
            m_fmt->nSamplesPerSec  = 44100;
            m_fmt->wBitsPerSample  = 16;
            m_fmt->nBlockAlign     = m_fmt->nChannels * m_fmt->wBitsPerSample / 8;
            m_fmt->nAvgBytesPerSec = m_fmt->nSamplesPerSec * m_fmt->nBlockAlign;
            hr = m_client->Initialize(AUDCLNT_SHAREMODE_SHARED,
                                      AUDCLNT_STREAMFLAGS_LOOPBACK,
                                      10000000, 0, m_fmt, nullptr);
            if (FAILED(hr)) return hr;
        }

        /* Check if WASAPI changed the format */
        WAVEFORMATEX* actual = nullptr;
        if (SUCCEEDED(m_client->GetMixFormat(&actual)) && actual) {
            if (actual->nSamplesPerSec != m_fmt->nSamplesPerSec ||
                actual->wBitsPerSample != m_fmt->wBitsPerSample) {
                CoTaskMemFree(m_fmt);
                m_fmt = actual;
            } else {
                CoTaskMemFree(actual);
            }
        }

        hr = m_client->GetService(__uuidof(IAudioCaptureClient), (void**)&m_capture);
        return hr;
    }

    HRESULT InitSystemWide() {
        m_processSpecific = false;

        ComPtr<IMMDeviceEnumerator> pEnum;
        HRESULT hr = CoCreateInstance(__uuidof(MMDeviceEnumerator), nullptr,
                                     CLSCTX_ALL, __uuidof(IMMDeviceEnumerator),
                                     (void**)&pEnum);
        if (FAILED(hr)) return hr;

        ComPtr<IMMDevice> pDev;
        hr = pEnum->GetDefaultAudioEndpoint(eRender, eConsole, &pDev);
        if (FAILED(hr)) return hr;

        hr = pDev->Activate(__uuidof(IAudioClient), CLSCTX_ALL, nullptr,
                            (void**)&m_client);
        if (FAILED(hr)) return hr;

        hr = m_client->GetMixFormat(&m_fmt);
        if (FAILED(hr)) return hr;

        hr = m_client->Initialize(AUDCLNT_SHAREMODE_SHARED,
                                  AUDCLNT_STREAMFLAGS_LOOPBACK |
                                  AUDCLNT_STREAMFLAGS_AUTOCONVERTPCM,
                                  10000000, 0, m_fmt, nullptr);
        if (FAILED(hr)) return hr;

        hr = m_client->GetService(__uuidof(IAudioCaptureClient), (void**)&m_capture);
        return hr;
    }
};

/* ── Python type wrapper ── */

typedef struct {
    PyObject_HEAD
    SystemLoopback* cap;
} LoopbackObject;

static void Loopback_dealloc(LoopbackObject* self) {
    if (self->cap) delete self->cap;
    Py_TYPE(self)->tp_free((PyObject*)self);
}

static PyObject* Loopback_new(PyTypeObject* type, PyObject*, PyObject*) {
    auto* self = (LoopbackObject*)type->tp_alloc(type, 0);
    if (self) self->cap = new SystemLoopback();
    return (PyObject*)self;
}

static int Loopback_init(LoopbackObject* self, PyObject*, PyObject*) {
    HRESULT hr = self->cap->Init();
    if (FAILED(hr)) {
        char msg[128];
        sprintf_s(msg, "WASAPI init failed (0x%08X). Need Win10 2004+.", hr);
        PyErr_SetString(PyExc_RuntimeError, msg);
        return -1;
    }
    return 0;
}

static PyObject* Loopback_start(LoopbackObject* self, PyObject*) {
    HRESULT hr = self->cap->Start();
    if (FAILED(hr)) {
        PyErr_Format(PyExc_RuntimeError, "Start failed: 0x%08X", hr);
        return nullptr;
    }
    Py_RETURN_NONE;
}

static PyObject* Loopback_stop(LoopbackObject* self, PyObject*) {
    self->cap->Stop();
    Py_RETURN_NONE;
}

static PyObject* Loopback_read(LoopbackObject* self, PyObject*) {
    return self->cap->Read();
}

static PyObject* Loopback_get_format(LoopbackObject* self, PyObject*) {
    return self->cap->GetFormat();
}

static PyObject* Loopback_is_excluding(LoopbackObject* self, PyObject*) {
    return PyBool_FromLong(self->cap->IsProcessSpecific() ? 1 : 0);
}

static PyMethodDef Loopback_methods[] = {
    {"start",        (PyCFunction)Loopback_start,        METH_NOARGS, "Start capture"},
    {"stop",         (PyCFunction)Loopback_stop,         METH_NOARGS, "Stop capture"},
    {"read",         (PyCFunction)Loopback_read,         METH_NOARGS, "Read captured PCM data"},
    {"get_format",   (PyCFunction)Loopback_get_format,   METH_NOARGS, "Get audio format dict"},
    {"is_excluding", (PyCFunction)Loopback_is_excluding, METH_NOARGS, "True if process-excluded capture"},
    {nullptr}
};

static PyTypeObject LoopbackType = {
    PyVarObject_HEAD_INIT(nullptr, 0)
    "_loopback.SystemLoopback",          /* tp_name */
    sizeof(LoopbackObject),              /* tp_basicsize */
    0,                                   /* tp_itemsize */
    (destructor)Loopback_dealloc,        /* tp_dealloc */
    0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    Py_TPFLAGS_DEFAULT,                  /* tp_flags */
    "WASAPI system loopback (excludes own process)",  /* tp_doc */
    0, 0, 0, 0, 0, 0,
    Loopback_methods,                    /* tp_methods */
    0, 0, 0, 0, 0, 0, 0,
    (initproc)Loopback_init,             /* tp_init */
    0,
    Loopback_new,                        /* tp_new */
};

static struct PyModuleDef module_def = {
    PyModuleDef_HEAD_INIT,
    "_loopback",
    "WASAPI system loopback capture (excludes own process audio)",
    -1, nullptr
};

PyMODINIT_FUNC PyInit__loopback(void) {
    if (PyType_Ready(&LoopbackType) < 0) return nullptr;
    PyObject* m = PyModule_Create(&module_def);
    if (!m) return nullptr;
    Py_INCREF(&LoopbackType);
    PyModule_AddObject(m, "SystemLoopback", (PyObject*)&LoopbackType);
    return m;
}