diff --git a/.github/native/README.md b/.github/native/README.md new file mode 100644 index 00000000..d28a33eb --- /dev/null +++ b/.github/native/README.md @@ -0,0 +1,28 @@ +# Recorder native SDK + +AIScan uses a two-stage build so ordinary full builds do not compile FFmpeg and x264. + +1. Maintainers run the `recorder-native-sdk` workflow after changing `versions.env` or the native build configuration. It builds the pinned sources, creates relocatable static SDK archives, writes SHA-256 sidecars, and publishes the assets to the versioned GitHub release. +2. Users, CI, and product release jobs run `fetch.sh`. It downloads the matching platform archive once, verifies it, and installs it below `.cache/record-native` before the Go/CGO link step. + +Supported bundles are `linux-amd64`, `linux-arm64`, and `windows-amd64`. FFmpeg and x264 are static; operating-system libraries remain external dependencies. The source builders use an explicit component allowlist (capture input, H.264 encoder, MP4 muxer, and file output only), and packaging rejects static-library sets larger than 16 MiB by default. + +Commands: + +```bash +# Consumer path (default for make full) +bash .github/native/fetch.sh linux amd64 + +# Maintainer/source path +bash .github/native/build-linux.sh +bash .github/native/package.sh linux amd64 dist/native +``` + +Environment overrides: + +- `AISCAN_RECORD_PREFIX`: SDK install/cache directory. +- `AISCAN_RECORD_NATIVE_URL`: release or mirror base URL containing the archive and `.sha256` sidecar. +- `AISCAN_RECORD_OFFLINE=1`: forbid downloads and require an already cached matching SDK. +- `AISCAN_RECORD_BUILD_FROM_SOURCE=1`: make `make full` or `build.sh -p full` use the pinned source builders instead of downloading an SDK. + +When native inputs or flags change, increment `RECORD_NATIVE_VERSION` and `RECORD_NATIVE_RELEASE` together before publishing. Do not replace an existing SDK version with incompatible contents. diff --git a/.github/native/build-linux.sh b/.github/native/build-linux.sh new file mode 100755 index 00000000..45af45f3 --- /dev/null +++ b/.github/native/build-linux.sh @@ -0,0 +1,85 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +source "${ROOT}/.github/native/versions.env" + +ARCH="$(go env GOARCH)" +PREFIX="${AISCAN_RECORD_PREFIX:-${ROOT}/.cache/record-native/linux-${ARCH}}" +SOURCE_ROOT="${AISCAN_RECORD_SOURCE:-${ROOT}/.cache/record-native/src}" +STAMP="${PREFIX}/.versions" +EXPECTED="source_bundle=${RECORD_NATIVE_VERSION} ffmpeg=${FFMPEG_COMMIT} x264=${X264_COMMIT}" + +emit_env() { + if [[ -n "${GITHUB_ENV:-}" ]]; then + echo "PKG_CONFIG_PATH=${PREFIX}/lib/pkgconfig" >> "${GITHUB_ENV}" + echo "PKG_CONFIG=${ROOT}/.github/native/pkg-config-static.sh" >> "${GITHUB_ENV}" + echo "CGO_CFLAGS=-I${PREFIX}/include" >> "${GITHUB_ENV}" + echo "CGO_LDFLAGS=-L${PREFIX}/lib" >> "${GITHUB_ENV}" + fi +} + +if [[ -f "${STAMP}" ]] && [[ "$(cat "${STAMP}")" == "${EXPECTED}" ]]; then + echo "record native dependencies already built at ${PREFIX}" + emit_env + exit 0 +fi + +mkdir -p "${SOURCE_ROOT}" "${PREFIX}" + +if [[ ! -d "${SOURCE_ROOT}/x264/.git" ]]; then + git clone "${X264_REPOSITORY}" "${SOURCE_ROOT}/x264" +fi +git -C "${SOURCE_ROOT}/x264" fetch --depth 1 origin "${X264_COMMIT}" +git -C "${SOURCE_ROOT}/x264" checkout --detach "${X264_COMMIT}" +( + cd "${SOURCE_ROOT}/x264" + make distclean >/dev/null 2>&1 || true + ./configure \ + --prefix="${PREFIX}" \ + --enable-static --enable-pic --disable-cli \ + --bit-depth=8 --chroma-format=420 \ + --disable-opencl --disable-interlaced + make -j"$(nproc)" + make install +) + +if [[ ! -d "${SOURCE_ROOT}/ffmpeg/.git" ]]; then + git clone "${FFMPEG_REPOSITORY}" "${SOURCE_ROOT}/ffmpeg" +fi +git -C "${SOURCE_ROOT}/ffmpeg" fetch --depth 1 origin "${FFMPEG_COMMIT}" +git -C "${SOURCE_ROOT}/ffmpeg" checkout --detach "${FFMPEG_COMMIT}" +( + cd "${SOURCE_ROOT}/ffmpeg" + make distclean >/dev/null 2>&1 || true + PKG_CONFIG_PATH="${PREFIX}/lib/pkgconfig" ./configure \ + --prefix="${PREFIX}" \ + --disable-shared --enable-static --enable-pic \ + --disable-programs --disable-doc --disable-debug --disable-network \ + --disable-autodetect --disable-everything \ + --enable-gpl --enable-libx264 \ + --enable-indev=xcbgrab --enable-decoder=rawvideo \ + --enable-encoder=libx264 --enable-muxer=mp4 \ + --enable-protocol=file --enable-swscale \ + --enable-libxcb --enable-libxcb-shm --enable-libxcb-shape --enable-libxcb-xfixes \ + --extra-cflags="-I${PREFIX}/include" \ + --extra-ldflags="-L${PREFIX}/lib" + bash "${ROOT}/.github/native/verify-ffmpeg-config.sh" linux config_components.h + make -j"$(nproc)" + make install +) + +mkdir -p "${PREFIX}/share/licenses/ffmpeg" "${PREFIX}/share/licenses/x264" +for license in COPYING.GPLv2 COPYING.GPLv3 LICENSE.md; do + if [[ -f "${SOURCE_ROOT}/ffmpeg/${license}" ]]; then + cp "${SOURCE_ROOT}/ffmpeg/${license}" "${PREFIX}/share/licenses/ffmpeg/" + fi +done +if [[ -f "${SOURCE_ROOT}/x264/COPYING" ]]; then + cp "${SOURCE_ROOT}/x264/COPYING" "${PREFIX}/share/licenses/x264/" +fi + +printf '%s' "${EXPECTED}" > "${STAMP}" +echo "record native dependencies built at ${PREFIX}" + +emit_env diff --git a/.github/native/build-windows.sh b/.github/native/build-windows.sh new file mode 100755 index 00000000..63f6cce3 --- /dev/null +++ b/.github/native/build-windows.sh @@ -0,0 +1,109 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +source "${ROOT}/.github/native/versions.env" + +# Calling usr/bin/bash directly does not enter the MINGW64 environment. Make +# the toolchain selection explicit so local builds and GitHub Actions agree. +export MSYSTEM=MINGW64 +export PATH="/mingw64/bin:/usr/bin:${PATH}" +if ! gcc -dumpmachine | grep -q 'mingw32$'; then + echo "a MinGW-w64 GCC toolchain is required" >&2 + exit 1 +fi + +ARCH="${GOARCH:-}" +if [[ -z "${ARCH}" ]] && command -v go >/dev/null 2>&1; then + ARCH="$(go env GOARCH)" +fi +if [[ -z "${ARCH}" ]]; then + case "$(uname -m)" in + x86_64|amd64) ARCH=amd64 ;; + *) echo "cannot determine recorder target architecture" >&2; exit 1 ;; + esac +fi +if [[ "${ARCH}" != "amd64" ]]; then + echo "record native Windows build only supports amd64" >&2 + exit 1 +fi + +PREFIX="${AISCAN_RECORD_PREFIX:-${ROOT}/.cache/record-native/windows-${ARCH}}" +SOURCE_ROOT="${AISCAN_RECORD_SOURCE:-${ROOT}/.cache/record-native/src}" +STAMP="${PREFIX}/.versions" +EXPECTED="source_bundle=${RECORD_NATIVE_VERSION} ffmpeg=${FFMPEG_COMMIT} x264=${X264_COMMIT}" + +emit_env() { + if [[ -n "${GITHUB_ENV:-}" ]]; then + PREFIX_WIN="$(cygpath -w "${PREFIX}")" + ROOT_WIN="$(cygpath -w "${ROOT}")" + echo "PKG_CONFIG_PATH=${PREFIX_WIN}\\lib\\pkgconfig" >> "${GITHUB_ENV}" + echo "PKG_CONFIG=${ROOT_WIN}\\.github\\native\\pkg-config-static.cmd" >> "${GITHUB_ENV}" + echo "CGO_CFLAGS=-I${PREFIX_WIN}\\include" >> "${GITHUB_ENV}" + echo "CGO_LDFLAGS=-L${PREFIX_WIN}\\lib -static -static-libgcc" >> "${GITHUB_ENV}" + fi +} + +if [[ -f "${STAMP}" ]] && [[ "$(cat "${STAMP}")" == "${EXPECTED}" ]]; then + echo "record native dependencies already built at ${PREFIX}" + emit_env + exit 0 +fi + +mkdir -p "${SOURCE_ROOT}" "${PREFIX}" + +if [[ ! -d "${SOURCE_ROOT}/x264/.git" ]]; then + git clone "${X264_REPOSITORY}" "${SOURCE_ROOT}/x264" +fi +git -C "${SOURCE_ROOT}/x264" fetch --depth 1 origin "${X264_COMMIT}" +git -C "${SOURCE_ROOT}/x264" checkout --detach "${X264_COMMIT}" +( + cd "${SOURCE_ROOT}/x264" + make distclean >/dev/null 2>&1 || true + ./configure \ + --prefix="${PREFIX}" \ + --host=x86_64-w64-mingw32 \ + --enable-static --disable-cli \ + --bit-depth=8 --chroma-format=420 \ + --disable-opencl --disable-interlaced + make -j"$(nproc)" + make install +) + +if [[ ! -d "${SOURCE_ROOT}/ffmpeg/.git" ]]; then + git clone "${FFMPEG_REPOSITORY}" "${SOURCE_ROOT}/ffmpeg" +fi +git -C "${SOURCE_ROOT}/ffmpeg" fetch --depth 1 origin "${FFMPEG_COMMIT}" +git -C "${SOURCE_ROOT}/ffmpeg" checkout --detach "${FFMPEG_COMMIT}" +( + cd "${SOURCE_ROOT}/ffmpeg" + make distclean >/dev/null 2>&1 || true + PKG_CONFIG_PATH="${PREFIX}/lib/pkgconfig" ./configure \ + --prefix="${PREFIX}" \ + --disable-shared --enable-static \ + --disable-programs --disable-doc --disable-debug --disable-network \ + --disable-autodetect --disable-everything \ + --enable-gpl --enable-libx264 \ + --enable-indev=gdigrab --enable-encoder=libx264 \ + --enable-muxer=mp4 --enable-protocol=file --enable-swscale \ + --extra-cflags="-I${PREFIX}/include" \ + --extra-ldflags="-L${PREFIX}/lib" + bash "${ROOT}/.github/native/verify-ffmpeg-config.sh" windows config_components.h + make -j"$(nproc)" + make install +) + +mkdir -p "${PREFIX}/share/licenses/ffmpeg" "${PREFIX}/share/licenses/x264" +for license in COPYING.GPLv2 COPYING.GPLv3 LICENSE.md; do + if [[ -f "${SOURCE_ROOT}/ffmpeg/${license}" ]]; then + cp "${SOURCE_ROOT}/ffmpeg/${license}" "${PREFIX}/share/licenses/ffmpeg/" + fi +done +if [[ -f "${SOURCE_ROOT}/x264/COPYING" ]]; then + cp "${SOURCE_ROOT}/x264/COPYING" "${PREFIX}/share/licenses/x264/" +fi + +printf '%s' "${EXPECTED}" > "${STAMP}" +echo "record native dependencies built at ${PREFIX}" + +emit_env diff --git a/.github/native/fetch.sh b/.github/native/fetch.sh new file mode 100755 index 00000000..1edc729e --- /dev/null +++ b/.github/native/fetch.sh @@ -0,0 +1,130 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +source "${ROOT}/.github/native/versions.env" + +detect_platform() { + case "$(uname -s)" in + Linux*) echo linux ;; + MINGW*|MSYS*|CYGWIN*) echo windows ;; + *) echo unsupported ;; + esac +} + +PLATFORM="${1:-$(detect_platform)}" +ARCH="${2:-$(go env GOARCH)}" +case "${PLATFORM}/${ARCH}" in + linux/amd64|linux/arm64|windows/amd64) ;; + *) + echo "no prebuilt recorder SDK for ${PLATFORM}/${ARCH}" >&2 + exit 1 + ;; +esac + +PREFIX="${AISCAN_RECORD_PREFIX:-${ROOT}/.cache/record-native/${PLATFORM}-${ARCH}}" +ARCHIVE="aiscan-record-native-${RECORD_NATIVE_VERSION}-${PLATFORM}-${ARCH}.tar.gz" +BASE_URL="${AISCAN_RECORD_NATIVE_URL:-https://github.com/${RECORD_NATIVE_REPOSITORY}/releases/download/${RECORD_NATIVE_RELEASE}}" +EXPECTED="bundle=${RECORD_NATIVE_VERSION} platform=${PLATFORM} arch=${ARCH} ffmpeg=${FFMPEG_COMMIT} x264=${X264_COMMIT}" +STAMP="${PREFIX}/.versions" + +emit_env() { + if [[ -z "${GITHUB_ENV:-}" ]]; then + return + fi + if [[ "${PLATFORM}" == "windows" ]] && command -v cygpath >/dev/null 2>&1; then + local prefix_native root_native + prefix_native="$(cygpath -w "${PREFIX}")" + root_native="$(cygpath -w "${ROOT}")" + echo "PKG_CONFIG_PATH=${prefix_native}\\lib\\pkgconfig" >> "${GITHUB_ENV}" + echo "PKG_CONFIG=${root_native}\\.github\\native\\pkg-config-static.cmd" >> "${GITHUB_ENV}" + echo "CGO_CFLAGS=-I${prefix_native}\\include" >> "${GITHUB_ENV}" + echo "CGO_LDFLAGS=-L${prefix_native}\\lib -static -static-libgcc" >> "${GITHUB_ENV}" + else + echo "PKG_CONFIG_PATH=${PREFIX}/lib/pkgconfig" >> "${GITHUB_ENV}" + echo "PKG_CONFIG=${ROOT}/.github/native/pkg-config-static.sh" >> "${GITHUB_ENV}" + echo "CGO_CFLAGS=-I${PREFIX}/include" >> "${GITHUB_ENV}" + echo "CGO_LDFLAGS=-L${PREFIX}/lib" >> "${GITHUB_ENV}" + fi +} + +if [[ -f "${STAMP}" ]] && [[ "$(cat "${STAMP}")" == "${EXPECTED}" ]]; then + echo "record native SDK already available at ${PREFIX}" + emit_env + exit 0 +fi + +if [[ "${AISCAN_RECORD_OFFLINE:-0}" == "1" ]]; then + echo "record native SDK is not cached at ${PREFIX} and offline mode is enabled" >&2 + exit 1 +fi + +for command_name in curl tar; do + if ! command -v "${command_name}" >/dev/null 2>&1; then + echo "${command_name} is required to download the recorder SDK" >&2 + exit 1 + fi +done + +case "${PREFIX}" in + ""|/|"${HOME:-__missing__}"|"${ROOT}") + echo "refusing unsafe recorder SDK prefix: ${PREFIX}" >&2 + exit 1 + ;; +esac + +TMP="$(mktemp -d)" +STAGE="${PREFIX}.tmp.$$" +BACKUP="${PREFIX}.old.$$" +cleanup() { + rm -rf "${TMP}" "${STAGE}" +} +trap cleanup EXIT + +echo "downloading recorder SDK ${RECORD_NATIVE_VERSION} for ${PLATFORM}/${ARCH}" +curl --fail --location --retry 3 --retry-delay 2 \ + "${BASE_URL}/${ARCHIVE}" -o "${TMP}/${ARCHIVE}" +curl --fail --location --retry 3 --retry-delay 2 \ + "${BASE_URL}/${ARCHIVE}.sha256" -o "${TMP}/${ARCHIVE}.sha256" + +if command -v sha256sum >/dev/null 2>&1; then + (cd "${TMP}" && sha256sum --check "${ARCHIVE}.sha256") +elif command -v shasum >/dev/null 2>&1; then + (cd "${TMP}" && shasum -a 256 --check "${ARCHIVE}.sha256") +else + echo "sha256sum or shasum is required to verify the recorder SDK" >&2 + exit 1 +fi + +mkdir -p "$(dirname "${PREFIX}")" +rm -rf "${STAGE}" "${BACKUP}" +mkdir -p "${STAGE}" +tar -xzf "${TMP}/${ARCHIVE}" -C "${STAGE}" + +if [[ ! -f "${STAGE}/.versions" ]] || [[ "$(cat "${STAGE}/.versions")" != "${EXPECTED}" ]]; then + echo "recorder SDK manifest does not match the requested version" >&2 + exit 1 +fi +for library in avcodec avdevice avfilter avformat avutil swresample swscale x264; do + if [[ ! -f "${STAGE}/lib/lib${library}.a" ]]; then + echo "recorder SDK archive is missing lib${library}.a" >&2 + exit 1 + fi +done +if [[ ! -d "${STAGE}/include/libavcodec" ]]; then + echo "recorder SDK archive is missing FFmpeg headers" >&2 + exit 1 +fi + +if [[ -e "${PREFIX}" ]]; then + mv "${PREFIX}" "${BACKUP}" +fi +if ! mv "${STAGE}" "${PREFIX}"; then + if [[ -e "${BACKUP}" ]]; then + mv "${BACKUP}" "${PREFIX}" + fi + exit 1 +fi +rm -rf "${BACKUP}" +echo "record native SDK installed at ${PREFIX}" +emit_env diff --git a/.github/native/package.sh b/.github/native/package.sh new file mode 100755 index 00000000..3508d21e --- /dev/null +++ b/.github/native/package.sh @@ -0,0 +1,78 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +source "${ROOT}/.github/native/versions.env" + +PLATFORM="${1:?usage: package.sh [output-dir]}" +ARCH="${2:?usage: package.sh [output-dir]}" +OUTPUT_DIR="${3:-${ROOT}/dist/native}" +PREFIX="${AISCAN_RECORD_PREFIX:-${ROOT}/.cache/record-native/${PLATFORM}-${ARCH}}" +SOURCE_STAMP="source_bundle=${RECORD_NATIVE_VERSION} ffmpeg=${FFMPEG_COMMIT} x264=${X264_COMMIT}" +BUNDLE_STAMP="bundle=${RECORD_NATIVE_VERSION} platform=${PLATFORM} arch=${ARCH} ffmpeg=${FFMPEG_COMMIT} x264=${X264_COMMIT}" +ARCHIVE="aiscan-record-native-${RECORD_NATIVE_VERSION}-${PLATFORM}-${ARCH}.tar.gz" +MAX_STATIC_LIB_BYTES="${AISCAN_RECORD_MAX_LIB_BYTES:-16777216}" + +case "${PLATFORM}/${ARCH}" in + linux/amd64|linux/arm64|windows/amd64) ;; + *) echo "unsupported recorder SDK target ${PLATFORM}/${ARCH}" >&2; exit 1 ;; +esac +if [[ ! -f "${PREFIX}/.versions" ]] || [[ "$(cat "${PREFIX}/.versions")" != "${SOURCE_STAMP}" ]]; then + echo "native dependencies at ${PREFIX} do not match versions.env" >&2 + exit 1 +fi + +static_lib_bytes=0 +while IFS= read -r -d '' library; do + bytes="$(wc -c < "${library}")" + static_lib_bytes=$((static_lib_bytes + bytes)) +done < <(find "${PREFIX}/lib" -maxdepth 1 -type f -name '*.a' -print0) +if (( static_lib_bytes > MAX_STATIC_LIB_BYTES )); then + echo "recorder static libraries are ${static_lib_bytes} bytes; budget is ${MAX_STATIC_LIB_BYTES}" >&2 + echo "the FFmpeg component allowlist may have regressed" >&2 + exit 1 +fi + +TMP="$(mktemp -d)" +cleanup() { rm -rf "${TMP}"; } +trap cleanup EXIT +STAGE="${TMP}/sdk" +mkdir -p "${STAGE}" "${OUTPUT_DIR}" +cp -R "${PREFIX}/include" "${PREFIX}/lib" "${STAGE}/" +if [[ -d "${PREFIX}/share/licenses" ]]; then + mkdir -p "${STAGE}/share" + cp -R "${PREFIX}/share/licenses" "${STAGE}/share/" +fi + +# Installed pkg-config files contain the maintainer's absolute build prefix. +# Make the SDK relocatable before publishing it. +if [[ -d "${STAGE}/lib/pkgconfig" ]]; then + while IFS= read -r -d '' pc; do + sed -i.bak 's|^prefix=.*|prefix=${pcfiledir}/../..|' "${pc}" + rm -f "${pc}.bak" + done < <(find "${STAGE}/lib/pkgconfig" -type f -name '*.pc' -print0) +fi + +printf '%s' "${BUNDLE_STAMP}" > "${STAGE}/.versions" +cat > "${STAGE}/README.txt" < "${OUTPUT_DIR}/${ARCHIVE}" + +if command -v sha256sum >/dev/null 2>&1; then + (cd "${OUTPUT_DIR}" && sha256sum "${ARCHIVE}" > "${ARCHIVE}.sha256") +else + digest="$(shasum -a 256 "${OUTPUT_DIR}/${ARCHIVE}" | awk '{print $1}')" + printf '%s %s\n' "${digest}" "${ARCHIVE}" > "${OUTPUT_DIR}/${ARCHIVE}.sha256" +fi +echo "packaged ${OUTPUT_DIR}/${ARCHIVE}" diff --git a/.github/native/pkg-config-static.cmd b/.github/native/pkg-config-static.cmd new file mode 100644 index 00000000..d428861f --- /dev/null +++ b/.github/native/pkg-config-static.cmd @@ -0,0 +1,2 @@ +@echo off +pkg-config --static %* diff --git a/.github/native/pkg-config-static.sh b/.github/native/pkg-config-static.sh new file mode 100755 index 00000000..e04b49e9 --- /dev/null +++ b/.github/native/pkg-config-static.sh @@ -0,0 +1,3 @@ +#!/usr/bin/env bash +set -euo pipefail +exec pkg-config --static "$@" diff --git a/.github/native/verify-ffmpeg-config.sh b/.github/native/verify-ffmpeg-config.sh new file mode 100755 index 00000000..f7661113 --- /dev/null +++ b/.github/native/verify-ffmpeg-config.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash +set -euo pipefail + +PLATFORM="${1:?usage: verify-ffmpeg-config.sh }" +CONFIG="${2:?usage: verify-ffmpeg-config.sh }" + +if [[ ! -f "${CONFIG}" ]]; then + echo "FFmpeg component config not found: ${CONFIG}" >&2 + exit 1 +fi + +enabled_components() { + local kind="$1" + sed -nE "s/^#define CONFIG_([A-Z0-9_]+)_${kind} 1$/\\1/p" "${CONFIG}" \ + | LC_ALL=C sort \ + | paste -sd, - +} + +expect_components() { + local kind="$1" + local expected="$2" + local actual + actual="$(enabled_components "${kind}")" + if [[ "${actual}" != "${expected}" ]]; then + echo "unexpected enabled FFmpeg ${kind,,} components" >&2 + echo "expected: ${expected:-}" >&2 + echo "actual: ${actual:-}" >&2 + exit 1 + fi +} + +case "${PLATFORM}" in + windows) + expect_components DECODER BMP + expect_components INDEV GDIGRAB + ;; + linux) + expect_components DECODER RAWVIDEO + expect_components INDEV XCBGRAB + ;; + *) + echo "unsupported recorder platform ${PLATFORM}" >&2 + exit 1 + ;; +esac + +expect_components ENCODER LIBX264 +expect_components MUXER MOV,MP4 +expect_components DEMUXER "" +expect_components PROTOCOL FILE +expect_components FILTER "" +expect_components OUTDEV "" +expect_components PARSER AC3 +expect_components BSF AAC_ADTSTOASC,VP9_SUPERFRAME + +echo "verified minimal FFmpeg component set for ${PLATFORM}" diff --git a/.github/native/versions.env b/.github/native/versions.env new file mode 100644 index 00000000..6e89c9ab --- /dev/null +++ b/.github/native/versions.env @@ -0,0 +1,8 @@ +RECORD_NATIVE_VERSION=ffmpeg-8.0.3-x264-0480cb05-2 +RECORD_NATIVE_RELEASE=record-native-ffmpeg-8.0.3-x264-0480cb05-2 +RECORD_NATIVE_REPOSITORY=chainreactors/aiscan +FFMPEG_TAG=n8.0.3 +FFMPEG_COMMIT=8ae0b34901ba60a802f183ee75a250a9fc3e09a5 +FFMPEG_REPOSITORY=https://github.com/FFmpeg/FFmpeg.git +X264_COMMIT=0480cb05fa188d37ae87e8f4fd8f1aea3711f7ee +X264_REPOSITORY=https://github.com/mirror/x264.git diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9aeeb4d5..48738557 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -310,6 +310,52 @@ jobs: CGO_ENABLED=1 go test -tags "full re2_cgo re2_static" -count=1 -timeout 5m \ ./cmd/aiscan ./pkg/web/api ./pkg/web/service + headless-record-replay-e2e: + runs-on: ubuntu-22.04 + needs: tidy + timeout-minutes: 10 + steps: + - name: Checkout + uses: actions/checkout@v6 + with: + fetch-depth: 0 + submodules: recursive + + - name: Set up Go + uses: actions/setup-go@v6 + with: + go-version-file: go.mod + cache: true + + - name: Set up Chrome + uses: browser-actions/setup-chrome@v2 + with: + chrome-version: stable + + - name: Verify Chrome discovery + run: chrome --version + + - name: Run headless action E2E + run: | + go test -tags "full re2_cgo re2_static" -count=1 -timeout 2m -v \ + -run '^(TestExecAIScanExtendedActions|TestExecAIScanHistoryActions)$' \ + ./pkg/headless + + - name: Run authenticated record and replay E2E + run: | + go test -tags "full re2_cgo re2_static" -count=1 -timeout 2m -v \ + -run '^TestE2E_RecordReplayAuthenticatedDashboard$' \ + ./tools/playwright + + - name: Run Katana browser reuse E2E + run: | + go test -count=1 -timeout 2m -v \ + -run '^TestE2EHeadlessReusesDiscoveredBrowser$' \ + ./tools/katana + go test -tags "full re2_cgo re2_static" -count=1 -timeout 2m -v \ + -run '^TestE2EKatanaDeepRendersAuthenticatedSPA$' \ + ./tools/scan + # ── Generated templates tests (depends on tidy) ─────────────── generated-test: @@ -531,6 +577,41 @@ jobs: if: runner.os == 'Windows' run: echo "C:/msys64/mingw64/bin" >> "$GITHUB_PATH" + - name: Install recorder SDK link dependencies on Linux + if: runner.os == 'Linux' + run: | + sudo apt-get update + sudo apt-get install -y build-essential nasm yasm pkg-config \ + libxcb1-dev libxcb-shm0-dev libxcb-shape0-dev libxcb-xfixes0-dev + + - name: Install recorder SDK link dependencies on Windows + if: runner.os == 'Windows' + run: | + C:/msys64/usr/bin/bash.exe -lc \ + "pacman -S --noconfirm --needed git diffutils make nasm yasm pkgconf mingw-w64-x86_64-toolchain" + + - name: Prepare static FFmpeg and x264 recorder SDK + if: runner.os != 'macOS' + run: | + if [[ "${RUNNER_OS}" == "Windows" ]]; then + C:/msys64/usr/bin/bash.exe -lc "cd '${GITHUB_WORKSPACE}'; bash .github/native/fetch.sh windows '${{ matrix.goarch }}' || bash .github/native/build-windows.sh" + root_native="$(cygpath -m "$(pwd)")" + prefix_native="${root_native}/.cache/record-native/windows-${{ matrix.goarch }}" + { + echo "PKG_CONFIG_PATH=${prefix_native}/lib/pkgconfig" + echo "PKG_CONFIG=${root_native}/.github/native/pkg-config-static.cmd" + echo "CGO_CFLAGS=-I${prefix_native}/include" + echo "CGO_LDFLAGS=-L${prefix_native}/lib -static -static-libgcc" + } >> "${GITHUB_ENV}" + else + chmod +x .github/native/pkg-config-static.sh + bash .github/native/fetch.sh linux '${{ matrix.goarch }}' || bash .github/native/build-linux.sh + fi + + - name: Verify recorder SDK link environment + if: runner.os != 'macOS' + run: pkg-config --modversion libavcodec + - name: Download embedded frontend uses: actions/download-artifact@v7 with: @@ -546,10 +627,28 @@ jobs: [[ "${{ matrix.goos }}" == "windows" ]] && suffix=".exe" CGO_ENABLED=1 GOOS="${{ matrix.goos }}" GOARCH="${{ matrix.goarch }}" \ go build -trimpath \ - -tags "forceposix emptytemplates noembed osusergo netgo full sqlite re2_cgo re2_static" \ + -tags "forceposix emptytemplates noembed osusergo netgo full sqlite record_ffmpeg re2_cgo re2_static" \ -ldflags "-s -w" -buildvcs=false \ -o "dist/full_${{ matrix.goos }}_${{ matrix.goarch }}${suffix}" ./cmd/aiscan + - name: Verify recorder libraries are statically linked + if: runner.os != 'macOS' + run: | + set -euo pipefail + binary="$(find dist -maxdepth 1 -type f -name 'full_*' -print -quit)" + test -n "${binary}" + if [[ "${RUNNER_OS}" == "Windows" ]]; then + if objdump -p "${binary}" | grep -Eiq 'DLL Name:.*(libav|x264|libwinpthread)'; then + echo "recorder library remained dynamically linked" >&2 + exit 1 + fi + else + if ldd "${binary}" | grep -Eiq '(libav|libx264)'; then + echo "recorder library remained dynamically linked" >&2 + exit 1 + fi + fi + - name: Upload full ${{ matrix.id }} binary uses: actions/upload-artifact@v7 with: diff --git a/.github/workflows/go-release.yml b/.github/workflows/go-release.yml index 827693d5..71901767 100644 --- a/.github/workflows/go-release.yml +++ b/.github/workflows/go-release.yml @@ -138,7 +138,7 @@ jobs: runner: ubuntu-22.04 main: ./cmd/aiscan binary: aiscan-full - tags: "forceposix emptytemplates noembed osusergo netgo full sqlite re2_cgo re2_static" + tags: "forceposix emptytemplates noembed osusergo netgo full sqlite record_ffmpeg re2_cgo re2_static" targets: "linux/amd64" cgo: "1" - id: aiscan-full-linux-arm64 @@ -146,7 +146,7 @@ jobs: runner: ubuntu-24.04-arm main: ./cmd/aiscan binary: aiscan-full - tags: "forceposix emptytemplates noembed osusergo netgo full sqlite re2_cgo re2_static" + tags: "forceposix emptytemplates noembed osusergo netgo full sqlite record_ffmpeg re2_cgo re2_static" targets: "linux/arm64" cgo: "1" - id: aiscan-full-darwin-amd64 @@ -154,7 +154,7 @@ jobs: runner: macos-15-intel main: ./cmd/aiscan binary: aiscan-full - tags: "forceposix emptytemplates noembed osusergo netgo full sqlite re2_cgo re2_static" + tags: "forceposix emptytemplates noembed osusergo netgo full sqlite record_ffmpeg re2_cgo re2_static" targets: "darwin/amd64" cgo: "1" - id: aiscan-full-darwin-arm64 @@ -170,7 +170,7 @@ jobs: runner: windows-2022 main: ./cmd/aiscan binary: aiscan-full - tags: "forceposix emptytemplates noembed osusergo netgo full sqlite re2_cgo re2_static" + tags: "forceposix emptytemplates noembed osusergo netgo full sqlite record_ffmpeg re2_cgo re2_static" targets: "windows/amd64" cgo: "1" @@ -196,6 +196,29 @@ jobs: if: runner.os == 'Windows' run: echo "C:/msys64/mingw64/bin" >> "$GITHUB_PATH" + - name: Install recorder SDK link dependencies on Linux + if: matrix.profile == 'full' && runner.os == 'Linux' + run: | + sudo apt-get update + sudo apt-get install -y build-essential pkg-config \ + libxcb1-dev libxcb-shm0-dev libxcb-shape0-dev libxcb-xfixes0-dev + + - name: Install recorder SDK link dependencies on Windows + if: matrix.profile == 'full' && runner.os == 'Windows' + run: | + C:/msys64/usr/bin/bash.exe -lc \ + "pacman -S --noconfirm --needed pkgconf mingw-w64-x86_64-toolchain" + + - name: Download static FFmpeg and x264 recorder SDK + if: matrix.profile == 'full' && runner.os != 'macOS' + run: | + if [[ "${RUNNER_OS}" == "Windows" ]]; then + C:/msys64/usr/bin/bash.exe -lc "cd '${GITHUB_WORKSPACE}'; bash .github/native/fetch.sh windows amd64" + else + chmod +x .github/native/pkg-config-static.sh + bash .github/native/fetch.sh linux "$(go env GOARCH)" + fi + - name: Download embedded frontend if: matrix.profile == 'full' uses: actions/download-artifact@v7 @@ -238,6 +261,24 @@ jobs: echo "=== Binaries ===" ls -lh "${OUTDIR}/" + - name: Verify recorder libraries are statically linked + if: matrix.profile == 'full' && runner.os != 'macOS' + run: | + set -euo pipefail + binary="$(find dist/build -maxdepth 1 -type f -name 'aiscan-full_*' -print -quit)" + test -n "${binary}" + if [[ "${RUNNER_OS}" == "Windows" ]]; then + if objdump -p "${binary}" | grep -Eiq 'DLL Name:.*(libav|x264|libwinpthread)'; then + echo "recorder library remained dynamically linked" >&2 + exit 1 + fi + else + if ldd "${binary}" | grep -Eiq '(libav|libx264)'; then + echo "recorder library remained dynamically linked" >&2 + exit 1 + fi + fi + - name: Upload artifacts uses: actions/upload-artifact@v7 with: diff --git a/.github/workflows/record-native.yml b/.github/workflows/record-native.yml new file mode 100644 index 00000000..be116a83 --- /dev/null +++ b/.github/workflows/record-native.yml @@ -0,0 +1,105 @@ +name: recorder-native-sdk + +on: + workflow_dispatch: + +permissions: + contents: write + +concurrency: + group: recorder-native-sdk + cancel-in-progress: false + +jobs: + build: + strategy: + fail-fast: false + matrix: + include: + - id: linux-amd64 + runner: ubuntu-22.04 + platform: linux + arch: amd64 + - id: linux-arm64 + runner: ubuntu-24.04-arm + platform: linux + arch: arm64 + - id: windows-amd64 + runner: windows-2022 + platform: windows + arch: amd64 + runs-on: ${{ matrix.runner }} + defaults: + run: + shell: bash + steps: + - name: Checkout + uses: actions/checkout@v6 + + - name: Set up Go + uses: actions/setup-go@v6 + with: + go-version-file: go.mod + cache: false + + - name: Install Linux source-build dependencies + if: runner.os == 'Linux' + run: | + sudo apt-get update + sudo apt-get install -y build-essential nasm yasm pkg-config \ + libxcb1-dev libxcb-shm0-dev libxcb-shape0-dev libxcb-xfixes0-dev + + - name: Install Windows source-build dependencies + if: runner.os == 'Windows' + run: | + C:/msys64/usr/bin/bash.exe -lc \ + "pacman -S --noconfirm --needed git diffutils make nasm yasm pkgconf mingw-w64-x86_64-toolchain" + + - name: Build and package Linux SDK + if: runner.os == 'Linux' + run: | + bash .github/native/build-linux.sh + bash .github/native/package.sh linux '${{ matrix.arch }}' dist/native + + - name: Build and package Windows SDK + if: runner.os == 'Windows' + run: | + C:/msys64/usr/bin/bash.exe -lc \ + "cd '${GITHUB_WORKSPACE}'; bash .github/native/build-windows.sh; bash .github/native/package.sh windows amd64 dist/native" + + - name: Upload SDK archive + uses: actions/upload-artifact@v7 + with: + name: recorder-native-${{ matrix.id }} + path: dist/native/* + if-no-files-found: error + retention-days: 7 + + publish: + needs: build + runs-on: ubuntu-22.04 + steps: + - name: Checkout + uses: actions/checkout@v6 + + - name: Download SDK archives + uses: actions/download-artifact@v7 + with: + pattern: recorder-native-* + path: dist/native + merge-multiple: true + + - name: Publish versioned SDK release + env: + GH_TOKEN: ${{ github.token }} + shell: bash + run: | + set -euo pipefail + source .github/native/versions.env + if gh release view "${RECORD_NATIVE_RELEASE}" >/dev/null 2>&1; then + gh release upload "${RECORD_NATIVE_RELEASE}" dist/native/* --clobber + else + gh release create "${RECORD_NATIVE_RELEASE}" dist/native/* \ + --title "AIScan recorder native SDK ${RECORD_NATIVE_VERSION}" \ + --notes "Prebuilt static FFmpeg ${FFMPEG_TAG} and x264 ${X264_COMMIT} SDKs used by AIScan full builds." + fi diff --git a/.gitignore b/.gitignore index 9d96fd59..878a0313 100644 --- a/.gitignore +++ b/.gitignore @@ -37,6 +37,7 @@ community.yaml /aiscan-deploy.yaml /*.log /.claude/ +/.cache/ # operator scan outputs dumped at repo root (screenshots, IP lists, app dumps, findings/reports) /*.png /*_ips.txt diff --git a/Makefile b/Makefile index eb5c93b8..bbd956b6 100644 --- a/Makefile +++ b/Makefile @@ -20,10 +20,28 @@ FULL_BIN ?= $(BIN_DIR)/aiscan-full$(EXE) # Standard/full match release artifacts. STANDARD_TAGS := forceposix emptytemplates noembed osusergo netgo -FULL_TAGS := forceposix emptytemplates noembed osusergo netgo full sqlite +FULL_TAGS := forceposix emptytemplates noembed osusergo netgo full sqlite record_ffmpeg BUILD_FLAGS := -trimpath -buildvcs=false -.PHONY: help prepare frontend proto-gen aop-gen standard full web-build web-run web all clean +UNAME_S := $(shell uname -s 2>/dev/null) +ifeq ($(OS),Windows_NT) +RECORD_PLATFORM := windows +else ifeq ($(UNAME_S),Linux) +RECORD_PLATFORM := linux +else +RECORD_PLATFORM := unsupported +endif +RECORD_PREFIX := $(if $(AISCAN_RECORD_PREFIX),$(AISCAN_RECORD_PREFIX),$(CURDIR)/.cache/record-native/$(RECORD_PLATFORM)-$(shell $(GO) env GOARCH)) +ifeq ($(RECORD_PLATFORM),windows) +RECORD_PKG_CONFIG := $(CURDIR)/.github/native/pkg-config-static.cmd +RECORD_EXTRA_LDFLAGS := -static -static-libgcc +else +RECORD_PKG_CONFIG := $(CURDIR)/.github/native/pkg-config-static.sh +RECORD_EXTRA_LDFLAGS := +endif +RECORD_BUILD_ENV := PKG_CONFIG="$(RECORD_PKG_CONFIG)" PKG_CONFIG_PATH="$(RECORD_PREFIX)/lib/pkgconfig" CGO_CFLAGS="-I$(RECORD_PREFIX)/include" CGO_LDFLAGS="-L$(RECORD_PREFIX)/lib $(RECORD_EXTRA_LDFLAGS)" + +.PHONY: help prepare frontend proto-gen aop-gen standard record-native record-native-source full web-build web-run web all clean help: @echo "AIScan build targets:" @@ -31,6 +49,8 @@ help: @echo " make full Build frontend, then build the full edition" @echo " make web Build the full edition and start the Web UI" @echo " make frontend Build only web/frontend into web/static" + @echo " make record-native Download the prebuilt FFmpeg/x264 recorder SDK" + @echo " make record-native-source Build the recorder SDK from pinned sources" @echo " make proto-gen Regenerate all AOP and AIScan protobuf bindings" @echo " make all Build the standard and full editions" @echo "" @@ -56,8 +76,26 @@ standard: prepare @echo "Built standard edition: $(STANDARD_BIN)" # The full binary embeds web/static, so frontend must finish first. -full: frontend prepare - CGO_ENABLED=1 $(GO) build $(BUILD_FLAGS) -ldflags "$(CGO_LDFLAGS)" -tags "$(FULL_TAGS)" -o "$(FULL_BIN)" ./cmd/aiscan +record-native: +ifeq ($(RECORD_PLATFORM),unsupported) + @echo "record native backend is not supported on this platform" +else + @if [ "$(AISCAN_RECORD_BUILD_FROM_SOURCE)" = "1" ]; then \ + bash ".github/native/build-$(RECORD_PLATFORM).sh"; \ + else \ + bash ".github/native/fetch.sh" "$(RECORD_PLATFORM)" "$(shell $(GO) env GOARCH)"; \ + fi +endif + +record-native-source: +ifeq ($(RECORD_PLATFORM),unsupported) + @echo "record native backend is not supported on this platform" +else + bash ".github/native/build-$(RECORD_PLATFORM).sh" +endif + +full: frontend record-native prepare + $(RECORD_BUILD_ENV) CGO_ENABLED=1 $(GO) build $(BUILD_FLAGS) -ldflags "$(CGO_LDFLAGS)" -tags "$(FULL_TAGS)" -o "$(FULL_BIN)" ./cmd/aiscan @echo "Built full edition: $(FULL_BIN)" web-build: full diff --git a/agent/hooks/hooks.go b/agent/hooks/hooks.go index 816fc143..9ee6f633 100644 --- a/agent/hooks/hooks.go +++ b/agent/hooks/hooks.go @@ -12,6 +12,7 @@ import ( "context" "errors" "fmt" + "runtime/debug" "sync" "sync/atomic" ) @@ -32,6 +33,8 @@ type HandlerError struct { Source string Kind Kind Err error + Panic any + Stack []byte } func (e *HandlerError) Error() string { @@ -289,9 +292,9 @@ func (p Point[E, R]) dispatch(ctx context.Context, r *Registry, entries []entry, } continue } - out, err := fn(ctx, ev) + out, err, panicValue, stack := invokeHandler(ctx, fn, ev) if err != nil { - he := &HandlerError{Source: e.source, Kind: p.Kind, Err: err} + he := &HandlerError{Source: e.source, Kind: p.Kind, Err: err, Panic: panicValue, Stack: stack} r.report(he) errs = append(errs, he) if p.OnError == FailClosed { @@ -308,3 +311,17 @@ func (p Point[E, R]) dispatch(ctx context.Context, r *Registry, entries []entry, } return acc, errors.Join(errs...) } + +func invokeHandler[E any, R any](ctx context.Context, fn func(context.Context, E) (R, error), ev E) (out R, err error, panicValue any, stack []byte) { + defer func() { + if recovered := recover(); recovered != nil { + var zero R + out = zero + err = errors.New("handler panicked") + panicValue = recovered + stack = debug.Stack() + } + }() + out, err = fn(ctx, ev) + return out, err, nil, nil +} diff --git a/agent/hooks/hooks_test.go b/agent/hooks/hooks_test.go index 6c072b7a..9f5ae816 100644 --- a/agent/hooks/hooks_test.go +++ b/agent/hooks/hooks_test.go @@ -135,6 +135,45 @@ func TestFailClosedShortCircuits(t *testing.T) { } } +func TestHandlerPanicIsAttributedAndReported(t *testing.T) { + r := New() + var reported *HandlerError + r.SetErrorSink(func(he *HandlerError) { reported = he }) + ToolCallHook.On(r, "plugin", func(context.Context, ToolCallEvent) (ToolCallResult, error) { + panic("boom") + }) + + _, err := ToolCallHook.Emit(context.Background(), r, ToolCallEvent{}) + if err == nil { + t.Fatal("err = nil, want handler panic") + } + if reported == nil || reported.Source != "plugin" || reported.Kind != "tool_call" { + t.Fatalf("reported = %+v", reported) + } + if reported.Panic != "boom" || len(reported.Stack) == 0 || strings.Contains(err.Error(), "boom") { + t.Fatalf("panic visibility = %+v, err = %v", reported, err) + } +} + +func TestContinueOnErrorContinuesAfterHandlerPanic(t *testing.T) { + r := New() + var secondRan bool + Context.On(r, "plugin", func(context.Context, ContextEvent) (ContextResult, error) { + panic("boom") + }) + Context.On(r, "core", func(context.Context, ContextEvent) (ContextResult, error) { + secondRan = true + return ContextResult{}, nil + }) + + if _, err := Context.Emit(context.Background(), r, ContextEvent{}); err == nil { + t.Fatal("err = nil, want handler panic") + } + if !secondRan { + t.Fatal("continue-on-error hook stopped after panic") + } +} + func TestContinueOnErrorCollectsAndKeepsGoing(t *testing.T) { r := New() first := errors.New("first") diff --git a/agent/loop.go b/agent/loop.go index e01b5385..70e1d650 100644 --- a/agent/loop.go +++ b/agent/loop.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "runtime/debug" "sort" "strings" "sync" @@ -463,7 +464,7 @@ func executeToolCalls(ctx context.Context, cfg Config, em *aopEmitter, assistant defer wg.Done() defer func() { <-sem }() slots[i].startedAt = time.Now() - slots[i].result = runToolCall(ctx, cfg, assistant.message, slots[i].tc, turn) + slots[i].result = runToolCallSafely(ctx, cfg, assistant.message, slots[i].tc, turn) }() } wg.Wait() @@ -515,6 +516,23 @@ type toolExecution struct { flow ToolFlowDecision } +func runToolCallSafely(ctx context.Context, cfg Config, assistantMsg *aop.Message, tc *aop.ToolCall, turn int) (execution toolExecution) { + defer func() { + if recovered := recover(); recovered != nil { + cfg.Logger.Errorf( + "tool call panic turn=%d name=%s call_id=%s session_id=%s panic=%v\n%s", + turn, tc.Name, tc.Id, cfg.SessionID, recovered, debug.Stack(), + ) + message := fmt.Sprintf("tool %s failed unexpectedly (call_id=%s)", tc.Name, tc.Id) + execution = toolExecution{ + result: message, rawResult: message, isError: true, + err: fmt.Errorf("tool call failed unexpectedly"), + } + } + }() + return runToolCall(ctx, cfg, assistantMsg, tc, turn) +} + func runToolCall(ctx context.Context, cfg Config, assistantMsg *aop.Message, tc *aop.ToolCall, turn int) toolExecution { startedAt := time.Now() toolCtx := tool.ContextWithInvocation(ctx, tool.Invocation{ @@ -542,7 +560,7 @@ func runToolCall(ctx context.Context, cfg Config, assistantMsg *aop.Message, tc if toolResult.Terminate { execution.flow = ToolFlowTerminate } - if tool.ResultHasImages(toolResult) || toolResult.Terminate { + if tool.ResultHasMedia(toolResult) || toolResult.Terminate { execution.fullResult = toolResult } } @@ -564,12 +582,10 @@ func (e toolExecution) eventContent() []*aop.Content { } for _, block := range e.fullResult.Output { media := block.GetMedia() - if media == nil || media.Kind != "image" || media.Resource == nil { + if media == nil || media.Resource == nil { continue } - if data := media.Resource.GetData(); len(data) > 0 { - content = append(content, aop.Image(media.Resource.MediaType, data)) - } + content = append(content, block) } return content } diff --git a/agent/loop_test.go b/agent/loop_test.go index fd39076c..19426f43 100644 --- a/agent/loop_test.go +++ b/agent/loop_test.go @@ -1,6 +1,7 @@ package agent import ( + "bytes" "context" "fmt" "reflect" @@ -18,6 +19,51 @@ import ( "github.com/chainreactors/aiscan/pkg/commands" ) +func TestParallelToolCallRecoversExtensionPanic(t *testing.T) { + tools := commands.NewRegistry() + tools.RegisterTool(&recordingTool{name: "first", output: "first ok"}) + tools.RegisterTool(&recordingTool{name: "second", output: "second ok"}) + var logs bytes.Buffer + cfg := Config{ + Tools: tools, + Logger: telemetry.NewLogger(telemetry.LogConfig{Debug: true, Output: &logs}), + BeforeToolCall: func(_ context.Context, call BeforeToolCallContext) (*BeforeToolCallResult, error) { + if call.ToolCall.Name == "first" { + panic("before boom") + } + return nil, nil + }, + }.init() + firstArgs, _ := aop.JSONValue(map[string]any{}) + secondArgs, _ := aop.JSONValue(map[string]any{}) + assistant := &assistantTurn{ + message: &aop.Message{Role: "assistant"}, + toolCalls: []*aop.ToolCall{ + {Id: "call-first", Name: "first", Arguments: firstArgs}, + {Id: "call-second", Name: "second", Arguments: secondArgs}, + }, + } + + batch, err := executeToolCalls(context.Background(), cfg, cfg.emitter, assistant, 1) + if err != nil { + t.Fatal(err) + } + if len(batch.messages) != 2 { + t.Fatalf("messages = %d", len(batch.messages)) + } + first := provider.MessageToolResult(batch.messages[0]) + second := provider.MessageToolResult(batch.messages[1]) + if first == nil || !first.IsError || !strings.Contains(tool.ResultText(first), "call-first") { + t.Fatalf("first result = %+v", first) + } + if second == nil || second.IsError || tool.ResultText(second) != "second ok" { + t.Fatalf("second result = %+v", second) + } + if got := logs.String(); !strings.Contains(got, "before boom") || !strings.Contains(got, "call-first") { + t.Fatalf("panic log = %s", got) + } +} + func TestRunEmitsTurnEndAfterToolResults(t *testing.T) { tools := commands.NewRegistry() tools.RegisterTool(&recordingTool{name: "echo", output: "tool output"}) diff --git a/aop/file/protocol.pb.go b/aop/file/protocol.pb.go index 67d15a6b..69fefa4b 100644 --- a/aop/file/protocol.pb.go +++ b/aop/file/protocol.pb.go @@ -22,8 +22,12 @@ const ( ) type ReadRequest struct { - state protoimpl.MessageState `protogen:"open.v1"` - Path string `protobuf:"bytes,1,opt,name=path,proto3" json:"path,omitempty"` + state protoimpl.MessageState `protogen:"open.v1"` + Path string `protobuf:"bytes,1,opt,name=path,proto3" json:"path,omitempty"` + // offset and limit enable bounded reads for large artifacts. A zero limit + // preserves the original whole-file behavior for older clients. + Offset int64 `protobuf:"varint,2,opt,name=offset,proto3" json:"offset,omitempty"` + Limit int32 `protobuf:"varint,3,opt,name=limit,proto3" json:"limit,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -65,6 +69,20 @@ func (x *ReadRequest) GetPath() string { return "" } +func (x *ReadRequest) GetOffset() int64 { + if x != nil { + return x.Offset + } + return 0 +} + +func (x *ReadRequest) GetLimit() int32 { + if x != nil { + return x.Limit + } + return 0 +} + type WriteRequest struct { state protoimpl.MessageState `protogen:"open.v1"` Path string `protobuf:"bytes,1,opt,name=path,proto3" json:"path,omitempty"` @@ -334,13 +352,17 @@ func (x *Entry) GetSize() int64 { } type Result struct { - state protoimpl.MessageState `protogen:"open.v1"` - Path string `protobuf:"bytes,1,opt,name=path,proto3" json:"path,omitempty"` - Filename string `protobuf:"bytes,2,opt,name=filename,proto3" json:"filename,omitempty"` - Size int64 `protobuf:"varint,3,opt,name=size,proto3" json:"size,omitempty"` - Data []byte `protobuf:"bytes,4,opt,name=data,proto3" json:"data,omitempty"` - Entries []*Entry `protobuf:"bytes,5,rep,name=entries,proto3" json:"entries,omitempty"` - MediaType string `protobuf:"bytes,6,opt,name=media_type,json=mediaType,proto3" json:"media_type,omitempty"` + state protoimpl.MessageState `protogen:"open.v1"` + Path string `protobuf:"bytes,1,opt,name=path,proto3" json:"path,omitempty"` + Filename string `protobuf:"bytes,2,opt,name=filename,proto3" json:"filename,omitempty"` + Size int64 `protobuf:"varint,3,opt,name=size,proto3" json:"size,omitempty"` + Data []byte `protobuf:"bytes,4,opt,name=data,proto3" json:"data,omitempty"` + Entries []*Entry `protobuf:"bytes,5,rep,name=entries,proto3" json:"entries,omitempty"` + MediaType string `protobuf:"bytes,6,opt,name=media_type,json=mediaType,proto3" json:"media_type,omitempty"` + // offset is the position of data within the file; size remains the total + // file size. eof marks the final chunk. + Offset int64 `protobuf:"varint,7,opt,name=offset,proto3" json:"offset,omitempty"` + Eof bool `protobuf:"varint,8,opt,name=eof,proto3" json:"eof,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -417,6 +439,20 @@ func (x *Result) GetMediaType() string { return "" } +func (x *Result) GetOffset() int64 { + if x != nil { + return x.Offset + } + return 0 +} + +func (x *Result) GetEof() bool { + if x != nil { + return x.Eof + } + return false +} + type ProtocolMessage struct { state protoimpl.MessageState `protogen:"open.v1"` // Types that are valid to be assigned to Message: @@ -567,9 +603,11 @@ var File_aop_file_protocol_proto protoreflect.FileDescriptor const file_aop_file_protocol_proto_rawDesc = "" + "\n" + - "\x17aop/file/protocol.proto\x12\baop.file\"!\n" + + "\x17aop/file/protocol.proto\x12\baop.file\"O\n" + "\vReadRequest\x12\x12\n" + - "\x04path\x18\x01 \x01(\tR\x04path\"6\n" + + "\x04path\x18\x01 \x01(\tR\x04path\x12\x16\n" + + "\x06offset\x18\x02 \x01(\x03R\x06offset\x12\x14\n" + + "\x05limit\x18\x03 \x01(\x05R\x05limit\"6\n" + "\fWriteRequest\x12\x12\n" + "\x04path\x18\x01 \x01(\tR\x04path\x12\x12\n" + "\x04data\x18\x02 \x01(\fR\x04data\"!\n" + @@ -587,7 +625,7 @@ const file_aop_file_protocol_proto_rawDesc = "" + "\x05Entry\x12\x12\n" + "\x04name\x18\x01 \x01(\tR\x04name\x12!\n" + "\fis_directory\x18\x02 \x01(\bR\visDirectory\x12\x12\n" + - "\x04size\x18\x03 \x01(\x03R\x04size\"\xaa\x01\n" + + "\x04size\x18\x03 \x01(\x03R\x04size\"\xd4\x01\n" + "\x06Result\x12\x12\n" + "\x04path\x18\x01 \x01(\tR\x04path\x12\x1a\n" + "\bfilename\x18\x02 \x01(\tR\bfilename\x12\x12\n" + @@ -595,7 +633,9 @@ const file_aop_file_protocol_proto_rawDesc = "" + "\x04data\x18\x04 \x01(\fR\x04data\x12)\n" + "\aentries\x18\x05 \x03(\v2\x0f.aop.file.EntryR\aentries\x12\x1d\n" + "\n" + - "media_type\x18\x06 \x01(\tR\tmediaType\"\x80\x03\n" + + "media_type\x18\x06 \x01(\tR\tmediaType\x12\x16\n" + + "\x06offset\x18\a \x01(\x03R\x06offset\x12\x10\n" + + "\x03eof\x18\b \x01(\bR\x03eof\"\x80\x03\n" + "\x0fProtocolMessage\x12:\n" + "\fread_request\x18\n" + " \x01(\v2\x15.aop.file.ReadRequestH\x00R\vreadRequest\x12=\n" + diff --git a/aop/helpers.go b/aop/helpers.go index fd1c39c0..2aafed02 100644 --- a/aop/helpers.go +++ b/aop/helpers.go @@ -80,11 +80,27 @@ func Reasoning(text string) *Content { } func Image(mediaType string, data []byte) *Content { + return MediaData("image", mediaType, "", data) +} + +func MediaData(kind, mediaType, filename string, data []byte) *Content { return &Content{Value: &Content_Media{Media: &MediaContent{ - Kind: "image", + Kind: kind, Resource: &Resource{ Source: &Resource_Data{Data: data}, MediaType: mediaType, + Filename: filename, + }, + }}} +} + +func MediaURI(kind, mediaType, filename, uri string) *Content { + return &Content{Value: &Content_Media{Media: &MediaContent{ + Kind: kind, + Resource: &Resource{ + Source: &Resource_Uri{Uri: uri}, + MediaType: mediaType, + Filename: filename, }, }}} } diff --git a/aop/helpers_test.go b/aop/helpers_test.go index 51d15004..e6d4abb6 100644 --- a/aop/helpers_test.go +++ b/aop/helpers_test.go @@ -43,3 +43,14 @@ func TestProviderFrameJSONAndBinaryRoundTrip(t *testing.T) { t.Fatalf("provider bytes changed") } } + +func TestMediaHelpersPreserveDataAndURI(t *testing.T) { + image := MediaData("image", "image/png", "shot.png", []byte("png")) + if media := image.GetMedia(); media.GetKind() != "image" || media.GetResource().GetFilename() != "shot.png" || string(media.GetResource().GetData()) != "png" { + t.Fatalf("image media = %+v", media) + } + video := MediaURI("video", "video/mp4", "capture.mp4", ".aiscan/record/capture.mp4") + if media := video.GetMedia(); media.GetKind() != "video" || media.GetResource().GetMediaType() != "video/mp4" || media.GetResource().GetUri() != ".aiscan/record/capture.mp4" { + t.Fatalf("video media = %+v", media) + } +} diff --git a/aop/tool/artifact.go b/aop/tool/artifact.go index dd11aebb..999b4545 100644 --- a/aop/tool/artifact.go +++ b/aop/tool/artifact.go @@ -5,6 +5,4 @@ const ( ArtifactKindWeb = "web" ArtifactKindWeakpass = "weakpass" ArtifactKindVuln = "vuln" - ArtifactKindSummary = "summary" - ArtifactKindError = "error" ) diff --git a/aop/tool/protocol.pb.go b/aop/tool/protocol.pb.go index cc730fa0..adb02d05 100644 --- a/aop/tool/protocol.pb.go +++ b/aop/tool/protocol.pb.go @@ -170,6 +170,7 @@ type Artifact struct { MediaType string `protobuf:"bytes,5,opt,name=media_type,json=mediaType,proto3" json:"media_type,omitempty"` Timestamp *timestamppb.Timestamp `protobuf:"bytes,6,opt,name=timestamp,proto3" json:"timestamp,omitempty"` CallId string `protobuf:"bytes,7,opt,name=call_id,json=callId,proto3" json:"call_id,omitempty"` + ResultId string `protobuf:"bytes,8,opt,name=result_id,json=resultId,proto3" json:"result_id,omitempty"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache } @@ -253,6 +254,123 @@ func (x *Artifact) GetCallId() string { return "" } +func (x *Artifact) GetResultId() string { + if x != nil { + return x.ResultId + } + return "" +} + +// Loot marks a scanner-native artifact as valuable without replacing or +// duplicating its evidence. result_id joins the marker back to Artifact. +type Loot struct { + state protoimpl.MessageState `protogen:"open.v1"` + ResultId string `protobuf:"bytes,1,opt,name=result_id,json=resultId,proto3" json:"result_id,omitempty"` + Tool string `protobuf:"bytes,2,opt,name=tool,proto3" json:"tool,omitempty"` + Kind string `protobuf:"bytes,3,opt,name=kind,proto3" json:"kind,omitempty"` + Target string `protobuf:"bytes,4,opt,name=target,proto3" json:"target,omitempty"` + Priority string `protobuf:"bytes,5,opt,name=priority,proto3" json:"priority,omitempty"` + Tags []string `protobuf:"bytes,6,rep,name=tags,proto3" json:"tags,omitempty"` + Description string `protobuf:"bytes,7,opt,name=description,proto3" json:"description,omitempty"` + VerificationStatus string `protobuf:"bytes,8,opt,name=verification_status,json=verificationStatus,proto3" json:"verification_status,omitempty"` + CallId string `protobuf:"bytes,9,opt,name=call_id,json=callId,proto3" json:"call_id,omitempty"` + unknownFields protoimpl.UnknownFields + sizeCache protoimpl.SizeCache +} + +func (x *Loot) Reset() { + *x = Loot{} + mi := &file_aop_tool_protocol_proto_msgTypes[3] + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + ms.StoreMessageInfo(mi) +} + +func (x *Loot) String() string { + return protoimpl.X.MessageStringOf(x) +} + +func (*Loot) ProtoMessage() {} + +func (x *Loot) ProtoReflect() protoreflect.Message { + mi := &file_aop_tool_protocol_proto_msgTypes[3] + if x != nil { + ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) + if ms.LoadMessageInfo() == nil { + ms.StoreMessageInfo(mi) + } + return ms + } + return mi.MessageOf(x) +} + +// Deprecated: Use Loot.ProtoReflect.Descriptor instead. +func (*Loot) Descriptor() ([]byte, []int) { + return file_aop_tool_protocol_proto_rawDescGZIP(), []int{3} +} + +func (x *Loot) GetResultId() string { + if x != nil { + return x.ResultId + } + return "" +} + +func (x *Loot) GetTool() string { + if x != nil { + return x.Tool + } + return "" +} + +func (x *Loot) GetKind() string { + if x != nil { + return x.Kind + } + return "" +} + +func (x *Loot) GetTarget() string { + if x != nil { + return x.Target + } + return "" +} + +func (x *Loot) GetPriority() string { + if x != nil { + return x.Priority + } + return "" +} + +func (x *Loot) GetTags() []string { + if x != nil { + return x.Tags + } + return nil +} + +func (x *Loot) GetDescription() string { + if x != nil { + return x.Description + } + return "" +} + +func (x *Loot) GetVerificationStatus() string { + if x != nil { + return x.VerificationStatus + } + return "" +} + +func (x *Loot) GetCallId() string { + if x != nil { + return x.CallId + } + return "" +} + type ProtocolMessage struct { state protoimpl.MessageState `protogen:"open.v1"` // Types that are valid to be assigned to Message: @@ -260,6 +378,7 @@ type ProtocolMessage struct { // *ProtocolMessage_Progress // *ProtocolMessage_Call // *ProtocolMessage_Artifact + // *ProtocolMessage_Loot Message isProtocolMessage_Message `protobuf_oneof:"message"` unknownFields protoimpl.UnknownFields sizeCache protoimpl.SizeCache @@ -267,7 +386,7 @@ type ProtocolMessage struct { func (x *ProtocolMessage) Reset() { *x = ProtocolMessage{} - mi := &file_aop_tool_protocol_proto_msgTypes[3] + mi := &file_aop_tool_protocol_proto_msgTypes[4] ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) ms.StoreMessageInfo(mi) } @@ -279,7 +398,7 @@ func (x *ProtocolMessage) String() string { func (*ProtocolMessage) ProtoMessage() {} func (x *ProtocolMessage) ProtoReflect() protoreflect.Message { - mi := &file_aop_tool_protocol_proto_msgTypes[3] + mi := &file_aop_tool_protocol_proto_msgTypes[4] if x != nil { ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x)) if ms.LoadMessageInfo() == nil { @@ -292,7 +411,7 @@ func (x *ProtocolMessage) ProtoReflect() protoreflect.Message { // Deprecated: Use ProtocolMessage.ProtoReflect.Descriptor instead. func (*ProtocolMessage) Descriptor() ([]byte, []int) { - return file_aop_tool_protocol_proto_rawDescGZIP(), []int{3} + return file_aop_tool_protocol_proto_rawDescGZIP(), []int{4} } func (x *ProtocolMessage) GetMessage() isProtocolMessage_Message { @@ -329,6 +448,15 @@ func (x *ProtocolMessage) GetArtifact() *Artifact { return nil } +func (x *ProtocolMessage) GetLoot() *Loot { + if x != nil { + if x, ok := x.Message.(*ProtocolMessage_Loot); ok { + return x.Loot + } + } + return nil +} + type isProtocolMessage_Message interface { isProtocolMessage_Message() } @@ -345,12 +473,18 @@ type ProtocolMessage_Artifact struct { Artifact *Artifact `protobuf:"bytes,12,opt,name=artifact,proto3,oneof"` } +type ProtocolMessage_Loot struct { + Loot *Loot `protobuf:"bytes,13,opt,name=loot,proto3,oneof"` +} + func (*ProtocolMessage_Progress) isProtocolMessage_Message() {} func (*ProtocolMessage_Call) isProtocolMessage_Message() {} func (*ProtocolMessage_Artifact) isProtocolMessage_Message() {} +func (*ProtocolMessage_Loot) isProtocolMessage_Message() {} + var File_aop_tool_protocol_proto protoreflect.FileDescriptor const file_aop_tool_protocol_proto_rawDesc = "" + @@ -366,7 +500,7 @@ const file_aop_tool_protocol_proto_rawDesc = "" + "\x06target\x18\x03 \x01(\tR\x06target\x128\n" + "\ttimestamp\x18\x05 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12\x12\n" + "\x04text\x18\x06 \x01(\tR\x04text\x12\x17\n" + - "\acall_id\x18\a \x01(\tR\x06callIdJ\x04\b\x02\x10\x03J\x04\b\x04\x10\x05\"\xd0\x01\n" + + "\acall_id\x18\a \x01(\tR\x06callIdJ\x04\b\x02\x10\x03J\x04\b\x04\x10\x05\"\xed\x01\n" + "\bArtifact\x12\x12\n" + "\x04tool\x18\x01 \x01(\tR\x04tool\x12\x12\n" + "\x04kind\x18\x02 \x01(\tR\x04kind\x12\x16\n" + @@ -375,12 +509,24 @@ const file_aop_tool_protocol_proto_rawDesc = "" + "\n" + "media_type\x18\x05 \x01(\tR\tmediaType\x128\n" + "\ttimestamp\x18\x06 \x01(\v2\x1a.google.protobuf.TimestampR\ttimestamp\x12\x17\n" + - "\acall_id\x18\a \x01(\tR\x06callId\"\xa6\x01\n" + + "\acall_id\x18\a \x01(\tR\x06callId\x12\x1b\n" + + "\tresult_id\x18\b \x01(\tR\bresultId\"\xff\x01\n" + + "\x04Loot\x12\x1b\n" + + "\tresult_id\x18\x01 \x01(\tR\bresultId\x12\x12\n" + + "\x04tool\x18\x02 \x01(\tR\x04tool\x12\x12\n" + + "\x04kind\x18\x03 \x01(\tR\x04kind\x12\x16\n" + + "\x06target\x18\x04 \x01(\tR\x06target\x12\x1a\n" + + "\bpriority\x18\x05 \x01(\tR\bpriority\x12\x12\n" + + "\x04tags\x18\x06 \x03(\tR\x04tags\x12 \n" + + "\vdescription\x18\a \x01(\tR\vdescription\x12/\n" + + "\x13verification_status\x18\b \x01(\tR\x12verificationStatus\x12\x17\n" + + "\acall_id\x18\t \x01(\tR\x06callId\"\xcc\x01\n" + "\x0fProtocolMessage\x120\n" + "\bprogress\x18\n" + " \x01(\v2\x12.aop.tool.ProgressH\x00R\bprogress\x12$\n" + "\x04call\x18\v \x01(\v2\x0e.aop.tool.CallH\x00R\x04call\x120\n" + - "\bartifact\x18\f \x01(\v2\x12.aop.tool.ArtifactH\x00R\bartifactB\t\n" + + "\bartifact\x18\f \x01(\v2\x12.aop.tool.ArtifactH\x00R\bartifact\x12$\n" + + "\x04loot\x18\r \x01(\v2\x0e.aop.tool.LootH\x00R\x04lootB\t\n" + "\amessageB/Z-github.com/chainreactors/aiscan/aop/tool;toolb\x06proto3" var ( @@ -395,27 +541,29 @@ func file_aop_tool_protocol_proto_rawDescGZIP() []byte { return file_aop_tool_protocol_proto_rawDescData } -var file_aop_tool_protocol_proto_msgTypes = make([]protoimpl.MessageInfo, 4) +var file_aop_tool_protocol_proto_msgTypes = make([]protoimpl.MessageInfo, 5) var file_aop_tool_protocol_proto_goTypes = []any{ (*Call)(nil), // 0: aop.tool.Call (*Progress)(nil), // 1: aop.tool.Progress (*Artifact)(nil), // 2: aop.tool.Artifact - (*ProtocolMessage)(nil), // 3: aop.tool.ProtocolMessage - (*aop.ToolCall)(nil), // 4: aop.ToolCall - (*timestamppb.Timestamp)(nil), // 5: google.protobuf.Timestamp + (*Loot)(nil), // 3: aop.tool.Loot + (*ProtocolMessage)(nil), // 4: aop.tool.ProtocolMessage + (*aop.ToolCall)(nil), // 5: aop.ToolCall + (*timestamppb.Timestamp)(nil), // 6: google.protobuf.Timestamp } var file_aop_tool_protocol_proto_depIdxs = []int32{ - 4, // 0: aop.tool.Call.call:type_name -> aop.ToolCall - 5, // 1: aop.tool.Progress.timestamp:type_name -> google.protobuf.Timestamp - 5, // 2: aop.tool.Artifact.timestamp:type_name -> google.protobuf.Timestamp + 5, // 0: aop.tool.Call.call:type_name -> aop.ToolCall + 6, // 1: aop.tool.Progress.timestamp:type_name -> google.protobuf.Timestamp + 6, // 2: aop.tool.Artifact.timestamp:type_name -> google.protobuf.Timestamp 1, // 3: aop.tool.ProtocolMessage.progress:type_name -> aop.tool.Progress 0, // 4: aop.tool.ProtocolMessage.call:type_name -> aop.tool.Call 2, // 5: aop.tool.ProtocolMessage.artifact:type_name -> aop.tool.Artifact - 6, // [6:6] is the sub-list for method output_type - 6, // [6:6] is the sub-list for method input_type - 6, // [6:6] is the sub-list for extension type_name - 6, // [6:6] is the sub-list for extension extendee - 0, // [0:6] is the sub-list for field type_name + 3, // 6: aop.tool.ProtocolMessage.loot:type_name -> aop.tool.Loot + 7, // [7:7] is the sub-list for method output_type + 7, // [7:7] is the sub-list for method input_type + 7, // [7:7] is the sub-list for extension type_name + 7, // [7:7] is the sub-list for extension extendee + 0, // [0:7] is the sub-list for field type_name } func init() { file_aop_tool_protocol_proto_init() } @@ -423,10 +571,11 @@ func file_aop_tool_protocol_proto_init() { if File_aop_tool_protocol_proto != nil { return } - file_aop_tool_protocol_proto_msgTypes[3].OneofWrappers = []any{ + file_aop_tool_protocol_proto_msgTypes[4].OneofWrappers = []any{ (*ProtocolMessage_Progress)(nil), (*ProtocolMessage_Call)(nil), (*ProtocolMessage_Artifact)(nil), + (*ProtocolMessage_Loot)(nil), } type x struct{} out := protoimpl.TypeBuilder{ @@ -434,7 +583,7 @@ func file_aop_tool_protocol_proto_init() { GoPackagePath: reflect.TypeOf(x{}).PkgPath(), RawDescriptor: unsafe.Slice(unsafe.StringData(file_aop_tool_protocol_proto_rawDesc), len(file_aop_tool_protocol_proto_rawDesc)), NumEnums: 0, - NumMessages: 4, + NumMessages: 5, NumExtensions: 0, NumServices: 0, }, diff --git a/build.sh b/build.sh index 7ca9e6ba..72a5484e 100755 --- a/build.sh +++ b/build.sh @@ -251,7 +251,7 @@ CGO_MODE=0 case "$PROFILE" in mini) ;; full) - EXTRA_TAGS="full${EXTRA_TAGS:+,$EXTRA_TAGS}" + EXTRA_TAGS="full,record_ffmpeg${EXTRA_TAGS:+,$EXTRA_TAGS}" BUILD_IOA=true AISCAN_BIN="aiscan-full" CGO_MODE=1 @@ -295,6 +295,29 @@ echo "cgo: $CGO_MODE" echo "output: $OUTPUT_DIR" echo "" +if [ "$PROFILE" = "full" ]; then + case "$HOST_OS" in + linux|windows) + if [ "${AISCAN_RECORD_BUILD_FROM_SOURCE:-0}" = "1" ]; then + bash ".github/native/build-${HOST_OS}.sh" + else + bash ".github/native/fetch.sh" "$HOST_OS" "$HOST_ARCH" + fi + RECORD_PREFIX="${AISCAN_RECORD_PREFIX:-$(pwd)/.cache/record-native/${HOST_OS}-${HOST_ARCH}}" + export PKG_CONFIG_PATH="${RECORD_PREFIX}/lib/pkgconfig" + export CGO_CFLAGS="-I${RECORD_PREFIX}/include" + export CGO_LDFLAGS="-L${RECORD_PREFIX}/lib" + if [ "$HOST_OS" = "windows" ]; then + export PKG_CONFIG="$(pwd)/.github/native/pkg-config-static.cmd" + export CGO_LDFLAGS="${CGO_LDFLAGS} -static -static-libgcc" + else + chmod +x ".github/native/pkg-config-static.sh" + export PKG_CONFIG="$(pwd)/.github/native/pkg-config-static.sh" + fi + ;; + esac +fi + # ─── 编译 ──────────────────────────────────────────────────────── mkdir -p "$OUTPUT_DIR" diff --git a/cmd/aiscan/imports_full.go b/cmd/aiscan/imports_full.go index cdf24a34..579a0c27 100644 --- a/cmd/aiscan/imports_full.go +++ b/cmd/aiscan/imports_full.go @@ -6,4 +6,5 @@ import ( _ "github.com/chainreactors/aiscan/tools/katana" _ "github.com/chainreactors/aiscan/tools/passive" _ "github.com/chainreactors/aiscan/tools/playwright" + _ "github.com/chainreactors/aiscan/tools/record" ) diff --git a/cmd/aiscan/imports_full_test.go b/cmd/aiscan/imports_full_test.go index 025208a7..22e18a45 100644 --- a/cmd/aiscan/imports_full_test.go +++ b/cmd/aiscan/imports_full_test.go @@ -3,15 +3,41 @@ package main import ( + "context" + "runtime" "slices" "testing" "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/core/telemetry" + "github.com/chainreactors/aiscan/pkg/runner" ) func TestFullCapabilitySet(t *testing.T) { want := []string{"arsenal", "browser", "core", "gogo", "ioa", "katana", "neutron", "passive", "proton", "proxy", "scan", "search", "spray", "zombie"} + if runtime.GOOS == "windows" || runtime.GOOS == "linux" { + want = append(want, "record") + slices.Sort(want) + } if got := capability.IDsSorted(); !slices.Equal(got, want) { t.Fatalf("full capabilities = %#v, want %#v", got, want) } } + +func TestFullRunnerBuildsDefaultRecordTool(t *testing.T) { + if runtime.GOOS != "windows" && runtime.GOOS != "linux" { + t.Skip("record is only linked on Windows and Linux") + } + app, err := runner.NewApp(context.Background(), runner.ApplicationConfig{ + Tools: runner.ToolConfig{BashTimeout: 1}, + Logger: telemetry.NopLogger(), + SkipEngines: true, + }) + if err != nil { + t.Fatal(err) + } + t.Cleanup(app.Close) + if _, ok := app.Commands.GetTool("record"); !ok { + t.Fatal("record tool is linked but was not assembled by the runner") + } +} diff --git a/cmd/aiscan/web_full.go b/cmd/aiscan/web_full.go index c37906ea..4166fe7a 100644 --- a/cmd/aiscan/web_full.go +++ b/cmd/aiscan/web_full.go @@ -8,6 +8,7 @@ import ( "io/fs" "net" "net/http" + "net/url" "os" "path" "path/filepath" @@ -157,16 +158,15 @@ func runWeb(ctx context.Context, option, explicitOption *cfg.Option, opts webCom // The hub's own agent comes online exactly like any node: an // `aiscan agent` dialed into this server over loopback WebSocket, // just in-process. The pool never sees a special "local" kind. - agentOption := *option - agentOption.ServerURL = "http://" + accessKey + "@" + listenAddr - if agentOption.IOANodeID == "" && agentOption.IOANodeName == "" { - agentOption.IOANodeName = "local" + agentOption, err := embeddedAgentOption(option, accessKey, listenAddr) + if err != nil { + return err } - go func() { + telemetry.SafeGo("embedded-agent", func() { if err := node.RunWebSocket(ctx, &agentOption, logger); err != nil && ctx.Err() == nil { logger.Warnf("embedded agent stopped: %s", err) } - }() + }) } if err := srv.Serve(listener); err != nil && err != http.ErrServerClosed { return err @@ -174,6 +174,24 @@ func runWeb(ctx context.Context, option, explicitOption *cfg.Option, opts webCom return nil } +func embeddedAgentOption(base *cfg.Option, accessKey, listenAddr string) (cfg.Option, error) { + var option cfg.Option + if base != nil { + option = *base + } + serverURL := &url.URL{Scheme: "http", Host: listenAddr} + serverURL.User = url.User(accessKey) + option.ServerURL = serverURL.String() + option.WebURL = option.ServerURL + if option.IOANodeID == "" && option.IOANodeName == "" { + option.IOANodeName = "local" + } + if err := cfg.ResolveAgentServerURLs(&option); err != nil { + return cfg.Option{}, fmt.Errorf("configure embedded agent: %w", err) + } + return option, nil +} + func wireWebApp(application *runner.App, ingestor webservice.ArtifactIngestor) { if application == nil || ingestor == nil || application.EventBus == nil { return diff --git a/cmd/aiscan/web_full_test.go b/cmd/aiscan/web_full_test.go index 9bc28896..bbac2624 100644 --- a/cmd/aiscan/web_full_test.go +++ b/cmd/aiscan/web_full_test.go @@ -90,6 +90,42 @@ func TestWireWebAppBindsRawArtifactsForReloadedApp(t *testing.T) { } } +func TestEmbeddedAgentOptionUsesSameOriginIOA(t *testing.T) { + base := &cfg.Option{IOAOptions: cfg.IOAOptions{Space: "case-1"}} + option, err := embeddedAgentOption(base, "promo-demo", "127.0.0.1:18080") + if err != nil { + t.Fatal(err) + } + if option.ServerURL != "http://promo-demo@127.0.0.1:18080" { + t.Fatalf("server URL = %q", option.ServerURL) + } + if option.IOAURL != "http://promo-demo@127.0.0.1:18080/ioa" { + t.Fatalf("IOA URL = %q, want embedded same-origin endpoint", option.IOAURL) + } + if option.IOANodeName != "local" || option.Space != "case-1" { + t.Fatalf("embedded identity = name %q space %q", option.IOANodeName, option.Space) + } + if base.ServerURL != "" || base.IOAURL != "" || base.IOANodeName != "" { + t.Fatalf("base option was mutated: %+v", base) + } +} + +func TestEmbeddedAgentOptionPreservesExplicitIOAAndNode(t *testing.T) { + base := &cfg.Option{ + IOAOptions: cfg.IOAOptions{ + IOAURL: "http://ioa-token@127.0.0.1:18765", + IOANodeName: "coordinator", + }, + } + option, err := embeddedAgentOption(base, "promo-demo", "127.0.0.1:18080") + if err != nil { + t.Fatal(err) + } + if option.IOAURL != base.IOAURL || option.IOANodeName != "coordinator" { + t.Fatalf("explicit IOA configuration was not preserved: %+v", option.IOAOptions) + } +} + type recordingArtifactIngestor struct { artifact *toolpb.Artifact } diff --git a/core/config/env.go b/core/config/env.go index 377ae8a1..b0821ded 100644 --- a/core/config/env.go +++ b/core/config/env.go @@ -41,10 +41,8 @@ func applyLLMEnvironment(option *Option, explicit Option, lookup envLookup) { if v := firstEnv(lookup, "AISCAN_PROVIDER"); v != "" && !providerExplicit { option.Provider = v } - - selectedProvider := selectedEnvProvider(option, lookup) - if option.Provider == "" && selectedProvider != "" && !providerExplicit { - option.Provider = selectedProvider + if option.Provider == "" && !providerExplicit { + option.Provider = firstEnv(lookup, "LLM_PROVIDER") } // AISCAN_BASE_URL is aiscan's own namespace: an intentional override that wins @@ -52,8 +50,15 @@ func applyLLMEnvironment(option *Option, explicit Option, lookup envLookup) { if strings.TrimSpace(explicit.BaseURL) == "" { if v := firstEnv(lookup, "AISCAN_BASE_URL"); v != "" { option.BaseURL = v + } else if strings.TrimSpace(option.BaseURL) == "" { + option.BaseURL = firstEnv(lookup, "LLM_BASE_URL") } } + + selectedProvider := selectedEnvProvider(option, lookup) + if option.Provider == "" && selectedProvider != "" && !providerExplicit { + option.Provider = selectedProvider + } // Provider-scoped base-URL envs (ANTHROPIC_BASE_URL, OPENAI_BASE_URL, …) are // commonly injected by the surrounding environment for *other* tools — e.g. // Claude-Code-style gateways export ANTHROPIC_BASE_URL. Treat them as a fallback @@ -73,6 +78,8 @@ func applyLLMEnvironment(option *Option, explicit Option, lookup envLookup) { if strings.TrimSpace(explicit.Model) == "" { if v := firstEnv(lookup, "AISCAN_MODEL"); v != "" { option.Model = v + } else if strings.TrimSpace(option.Model) == "" { + option.Model = firstEnv(lookup, "LLM_MODEL") } } // Provider-scoped model envs (ANTHROPIC_MODEL, OPENAI_MODEL, …) are commonly @@ -91,6 +98,8 @@ func applyLLMEnvironment(option *Option, explicit Option, lookup envLookup) { if strings.TrimSpace(explicit.APIKey) == "" { if v := firstEnv(lookup, "AISCAN_API_KEY"); v != "" { option.APIKey = v + } else if strings.TrimSpace(option.APIKey) == "" { + option.APIKey = firstEnv(lookup, "LLM_API_KEY") } } // Provider-scoped key envs (ANTHROPIC_API_KEY, OPENAI_API_KEY) are commonly diff --git a/core/config/loader_test.go b/core/config/loader_test.go index d7bfc49e..bf11e813 100644 --- a/core/config/loader_test.go +++ b/core/config/loader_test.go @@ -587,6 +587,27 @@ llm: }) } +func TestApplyEnvironmentUsesSharedLLMConfiguration(t *testing.T) { + values := map[string]string{ + "LLM_BASE_URL": "https://api.deepseek.com", + "LLM_API_KEY": "shared-key", + "LLM_MODEL": "deepseek-v4-flash", + } + lookup := func(name string) (string, bool) { + value, ok := values[name] + return value, ok + } + + option := Option{} + applyEnvironment(&option, Option{}, lookup) + if err := normalizeProviderOptions(&option); err != nil { + t.Fatal(err) + } + if option.Provider != "openai" || option.BaseURL != values["LLM_BASE_URL"] || option.APIKey != values["LLM_API_KEY"] || option.Model != values["LLM_MODEL"] { + t.Fatalf("shared LLM configuration not applied: %#v", option.LLMOptions) + } +} + func TestResolveRuntimeConfigUsesAnthropicEnvironment(t *testing.T) { t.Setenv("ANTHROPIC_BASE_URL", "https://anthropic-proxy.example/v1") t.Setenv("ANTHROPIC_MODEL", "claude-env") diff --git a/core/config/options.go b/core/config/options.go index fcf88abc..3655f2a0 100644 --- a/core/config/options.go +++ b/core/config/options.go @@ -76,7 +76,7 @@ type AgentOptions struct { Prompt string `short:"p" long:"prompt" description:"Natural language task or existing file path for the agent"` Inputs []string `short:"i" long:"input" description:"Target input: IP, URL, IP:port, or CIDR. Can specify multiple"` Skills []string `short:"s" long:"skill" description:"Skill to apply (name or file path). Can specify multiple"` - Tools []string `short:"t" long:"tools" config:"tools" description:"Optional tool groups to enable (search, browser). Arsenal is always loaded"` + Tools []string `short:"t" long:"tools" config:"tools" description:"Optional tool groups to enable. Arsenal is always loaded"` TaskFile string `long:"task-file" description:"File containing task description"` Heartbeat int `long:"heartbeat" description:"Heartbeat interval in minutes: periodically wake the agent to review context (0 disables)" default:"0"` Timeout int `long:"timeout" config:"timeout" description:"Overall timeout in seconds" default:"3600"` diff --git a/core/deps/deps_test.go b/core/deps/deps_test.go index 48691ef9..690dcc0a 100644 --- a/core/deps/deps_test.go +++ b/core/deps/deps_test.go @@ -459,9 +459,9 @@ func TestBuildProfilesUseExpectedCGOModes(t *testing.T) { makefile := readRepositoryFile(t, root, "Makefile") for _, required := range []string{ "standard: prepare\n\tCGO_ENABLED=0 $(GO) build", - "full: frontend prepare\n\tCGO_ENABLED=1 $(GO) build", + "full: frontend record-native prepare\n\t$(RECORD_BUILD_ENV) CGO_ENABLED=1 $(GO) build", "STANDARD_TAGS := forceposix emptytemplates noembed osusergo netgo", - "FULL_TAGS := forceposix emptytemplates noembed osusergo netgo full sqlite", + "FULL_TAGS := forceposix emptytemplates noembed osusergo netgo full sqlite record_ffmpeg", } { if !strings.Contains(makefile, required) { t.Errorf("Makefile missing build profile contract %q", required) diff --git a/core/output/types.go b/core/output/types.go index e9495808..0b8894b7 100644 --- a/core/output/types.go +++ b/core/output/types.go @@ -9,11 +9,22 @@ import ( // ScanResult is private collector state. Scanner-native records leave a node // only as canonical aop.tool.Artifact messages. type ScanResult struct { - Summary Summary `json:"summary"` - GOGO []*parsers.GOGOResult `json:"gogo,omitempty"` - Spray []*parsers.SprayResult `json:"spray,omitempty"` - Loots []Loot `json:"loots,omitempty"` - Errors []Error `json:"errors,omitempty"` + Summary Summary `json:"summary"` + GOGO []*parsers.GOGOResult `json:"gogo,omitempty"` + Spray []*parsers.SprayResult `json:"spray,omitempty"` + Artifacts []ArtifactResult `json:"artifacts,omitempty"` + Loots []Loot `json:"loots,omitempty"` + Errors []Error `json:"errors,omitempty"` +} + +// ArtifactResult keeps the scanner-native result paired with a Loot marker. +// Data is serialized directly into aop.tool.Artifact without reshaping. +type ArtifactResult struct { + ResultID string `json:"result_id"` + Tool string `json:"tool"` + Kind string `json:"kind"` + Target string `json:"target"` + Data any `json:"data"` } type Summary struct { diff --git a/core/tool/result.go b/core/tool/result.go index 150afa14..8213e7bc 100644 --- a/core/tool/result.go +++ b/core/tool/result.go @@ -34,6 +34,18 @@ func ResultHasImages(r *Result) bool { return false } +func ResultHasMedia(r *Result) bool { + if r == nil { + return false + } + for _, block := range r.Output { + if block.GetMedia() != nil { + return true + } + } + return false +} + func TextResult(s string) *Result { return &Result{Output: []*aop.Content{aop.Text(s)}} } diff --git a/docs/record.md b/docs/record.md new file mode 100644 index 00000000..2ecd0238 --- /dev/null +++ b/docs/record.md @@ -0,0 +1,67 @@ +# record — desktop and window capture + +`record` is a native full-build tool for Windows and Linux X11. It captures PNG screenshots and H.264/MP4 recordings from the desktop or a visible application window. + +Examples: + +```json +{"action":"screenshot"} +{"action":"screenshot","target":"window","pid":1234} +{"action":"record","target":"window","window_handle":"0x12345","duration_seconds":10} +{"action":"start","target":"desktop","fps":30} +{"action":"stop","recording_id":""} +{"action":"status"} +``` + +Windows window targets use an `HWND`; Linux uses an X11 Window ID. Handles are strings and accept decimal or `0x` hexadecimal notation. A PID resolves to the largest visible, non-minimized top-level window owned by that process. + +Defaults: + +- Desktop target, 30 FPS, mouse cursor included. +- Screenshots are PNG; recordings are H.264/libx264 in MP4. +- Outputs are written below `.aiscan/record/` unless `output` is specified. +- At most four recordings run concurrently. Set `AISCAN_RECORD_MAX_CONCURRENT` to a value from 1 to 16 to change the limit. + +Media transport uses the existing AOP media and file namespaces. Screenshot +previews are returned as bounded inline `Content.media` data. Completed videos +are returned as `Content.media` with a task-relative `Resource.uri`; consumers +read the underlying MP4 through chunked `aop.file` requests. When a tool +invocation supplies a work directory, the default output is +`/.aiscan/record/`, so remote runners can expose the URI without +leaking or depending on a machine-global data path. + +Limitations: + +- Video only; microphone and system audio are not captured. +- Wayland is not supported. Use an X11 session. +- The window must be visible and non-minimized. Capture size is fixed when recording starts; closing, minimizing, or shrinking the window can terminate the recording. +- The native backend is present in official Windows/Linux full builds. Custom builds require CGO and a supported C toolchain. `make full` downloads the pinned, prebuilt FFmpeg/x264 SDK automatically. + +The full build statically links a feature-minimal FFmpeg and x264. The SDK only enables the platform capture input, its raw/BMP decoder, libx264, the MP4 muxer, file output, and pixel conversion. It is not a general-purpose FFmpeg build. Windows system DLLs and Linux system libraries such as glibc/X11 remain platform dependencies. + +## Two-stage native build + +Normal users should use an official `aiscan-full` archive; recording works without installing FFmpeg or x264. Developers building from source use: + +```bash +make full +``` + +The `record-native` prerequisite downloads a versioned SDK into `.cache/record-native/-`, verifies its SHA-256 sidecar and manifest, then links it into the full binary. Supported SDK targets are Linux amd64/arm64 and Windows amd64. Linux source builds still need a C compiler, `pkg-config`, and XCB development packages; Windows source builds need MinGW-w64 and `pkgconf`. + +Maintainers build the SDK from the pinned commits separately: + +```bash +make record-native-source +bash .github/native/package.sh linux "$(go env GOARCH)" dist/native +``` + +The `recorder-native-sdk` GitHub Actions workflow performs that source-build/package phase for every supported target and publishes the archives under the release tag declared in `.github/native/versions.env`. Normal CI, release builds, `make full`, and `build.sh -p full` only consume those archives. Set `AISCAN_RECORD_BUILD_FROM_SOURCE=1` when invoking `make full` or `build.sh -p full` to opt into the slow source-build fallback. `AISCAN_RECORD_PREFIX` changes the SDK cache/install directory, and `AISCAN_RECORD_NATIVE_URL` can point downloads at an internal mirror. + +The source build verifies an exact FFmpeg component allowlist, and packaging rejects static libraries above a 16 MiB budget unless `AISCAN_RECORD_MAX_LIB_BYTES` explicitly overrides it. This prevents an FFmpeg upgrade or configure change from silently restoring all default codecs and adding tens of megabytes to `aiscan-full`. + +Native smoke tests are opt-in because they require an interactive desktop/X11 session: + +```bash +go test -tags "record_ffmpeg record_integration" ./tools/record +``` diff --git a/go.mod b/go.mod index 9090627f..d987d492 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( connectrpc.com/connect v1.20.0 github.com/Microsoft/go-winio v0.6.2 github.com/alecthomas/chroma/v2 v2.14.0 + github.com/asticode/go-astiav v0.41.0 github.com/carapace-sh/carapace v1.11.6 github.com/chainreactors/crtm v0.0.3-0.20260618163257-073207497076 github.com/chainreactors/fingers v1.2.2-0.20260714063144-070758342f45 @@ -62,6 +63,7 @@ require ( ) require ( + github.com/asticode/go-astikit v0.42.0 // indirect github.com/jinzhu/inflection v1.0.0 // indirect github.com/puzpuzpuz/xsync/v3 v3.5.1 // indirect github.com/tmthrgd/go-hex v0.0.0-20190904060850-447a3041c3bc // indirect diff --git a/go.sum b/go.sum index c5926678..260c6d5f 100644 --- a/go.sum +++ b/go.sum @@ -129,6 +129,10 @@ github.com/armon/go-radix v0.0.0-20180808171621-7fddfc383310/go.mod h1:ufUuZ+zHj github.com/armon/go-radix v1.0.0/go.mod h1:ufUuZ+zHj4x4TnLV4JWEpy2hxWSpsRywHrMgIH9cCH8= github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2 h1:DklsrG3dyBCFEj5IhUbnKptjxatkF07cF2ak3yi77so= github.com/asaskevich/govalidator v0.0.0-20230301143203-a9d515a09cc2/go.mod h1:WaHUgvxTVq04UNunO+XhnAqY/wQc+bxr74GqbsZ/Jqw= +github.com/asticode/go-astiav v0.41.0 h1:HZCQ71lPqRQHHIQ5cShrHEq9HtX0NiWIHm7FZ2Jk+x8= +github.com/asticode/go-astiav v0.41.0/go.mod h1:GI0pHw6K2/pl/o8upCtT49P/q4KCwhv/8nGLlCsZLdA= +github.com/asticode/go-astikit v0.42.0 h1:pnir/2KLUSr0527Tv908iAH6EGYYrYta132vvjXsH5w= +github.com/asticode/go-astikit v0.42.0/go.mod h1:h4ly7idim1tNhaVkdVBeXQZEE3L0xblP7fCWbgwipF0= github.com/atotto/clipboard v0.1.4 h1:EH0zSVneZPSuFR11BlR9YppQTVDbh5+16AmcJi4g1z4= github.com/atotto/clipboard v0.1.4/go.mod h1:ZY9tmq7sm5xIbd9bOK4onWV4S6X0u6GY7Vn0Yu86PYI= github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k= diff --git a/pkg/browser/discovery.go b/pkg/browser/discovery.go new file mode 100644 index 00000000..e98e391c --- /dev/null +++ b/pkg/browser/discovery.go @@ -0,0 +1,59 @@ +// Package browser centralizes browser binary discovery for AIScan's browser-backed engines. +package browser + +import ( + "fmt" + "os" + "os/exec" + "strings" + + "github.com/go-rod/rod/lib/launcher" +) + +const ( + // PathEnv explicitly selects the Chrome-compatible browser binary used by AIScan. + PathEnv = "AISCAN_BROWSER_PATH" +) + +// Source identifies how a browser binary was selected. +type Source string + +const ( + SourceEnvironment Source = "environment" + SourceSystem Source = "system" +) + +// Binary describes a discovered Chrome-compatible browser executable. +type Binary struct { + Path string + Source Source +} + +// Discover resolves the browser shared by Playwright, nuclei headless, and Katana. +// An explicit AISCAN_BROWSER_PATH is authoritative. If neither it nor a system +// browser is available, an empty result lets Rod use its cached/download fallback. +func Discover() (Binary, error) { + configured, configuredSet := os.LookupEnv(PathEnv) + return discover(configured, configuredSet, exec.LookPath, launcher.LookPath) +} + +func discover( + configured string, + configuredSet bool, + resolve func(string) (string, error), + findSystem func() (string, bool), +) (Binary, error) { + configured = strings.TrimSpace(configured) + if configuredSet && configured != "" { + path, err := resolve(configured) + if err != nil { + return Binary{}, fmt.Errorf("%s=%q does not resolve to an executable browser: %w", PathEnv, configured, err) + } + return Binary{Path: path, Source: SourceEnvironment}, nil + } + + if path, ok := findSystem(); ok && path != "" { + return Binary{Path: path, Source: SourceSystem}, nil + } + return Binary{}, nil +} diff --git a/pkg/browser/discovery_test.go b/pkg/browser/discovery_test.go new file mode 100644 index 00000000..f5f6d733 --- /dev/null +++ b/pkg/browser/discovery_test.go @@ -0,0 +1,90 @@ +package browser + +import ( + "errors" + "strings" + "testing" +) + +func TestDiscoverPriority(t *testing.T) { + tests := []struct { + name string + configured string + configuredSet bool + resolvePath string + resolveErr error + systemPath string + systemFound bool + want Binary + wantErr bool + }{ + { + name: "environment overrides system browser", + configured: " /opt/aiscan/chrome ", + configuredSet: true, + resolvePath: "/opt/aiscan/chrome", + systemPath: "/usr/bin/chrome", + systemFound: true, + want: Binary{Path: "/opt/aiscan/chrome", Source: SourceEnvironment}, + }, + { + name: "invalid environment is an error", + configured: "/missing/chrome", + configuredSet: true, + resolveErr: errors.New("not found"), + systemPath: "/usr/bin/chrome", + systemFound: true, + wantErr: true, + }, + { + name: "system browser is automatic fallback", + systemPath: "/usr/bin/chromium", + systemFound: true, + want: Binary{Path: "/usr/bin/chromium", Source: SourceSystem}, + }, + { + name: "blank environment still allows system discovery", + configured: " ", + configuredSet: true, + systemPath: "/usr/bin/edge", + systemFound: true, + want: Binary{Path: "/usr/bin/edge", Source: SourceSystem}, + }, + { + name: "empty result preserves Rod fallback", + want: Binary{}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + resolveCalls := 0 + findSystemCalls := 0 + got, err := discover( + tt.configured, + tt.configuredSet, + func(path string) (string, error) { + resolveCalls++ + return tt.resolvePath, tt.resolveErr + }, + func() (string, bool) { + findSystemCalls++ + return tt.systemPath, tt.systemFound + }, + ) + if (err != nil) != tt.wantErr { + t.Fatalf("discover error = %v, wantErr %v", err, tt.wantErr) + } + if got != tt.want { + t.Fatalf("discover = %#v, want %#v", got, tt.want) + } + explicit := tt.configuredSet && strings.TrimSpace(tt.configured) != "" + if explicit && resolveCalls != 1 { + t.Fatalf("resolve calls = %d, want 1", resolveCalls) + } + if explicit && findSystemCalls != 0 { + t.Fatalf("system discovery called %d times after explicit configuration", findSystemCalls) + } + }) + } +} diff --git a/pkg/commands/bash_test.go b/pkg/commands/bash_test.go index 0a990f72..be5481c8 100644 --- a/pkg/commands/bash_test.go +++ b/pkg/commands/bash_test.go @@ -719,17 +719,23 @@ func TestPseudoFlagWithPipeChar(t *testing.T) { func TestExecuteTool_RecoversPanic(t *testing.T) { reg := NewRegistry() + var logs bytes.Buffer + reg.SetLogger(telemetry.NewLogger(telemetry.LogConfig{Debug: true, Output: &logs})) reg.RegisterTool(&panicTool{msg: "boom"}) - result, err := reg.ExecuteTool(context.Background(), "panic_tool", "{}") + ctx := tool.ContextWithInvocation(context.Background(), tool.Invocation{CallID: "call-panic"}) + result, err := reg.ExecuteTool(ctx, "panic_tool", "{}") if err == nil { t.Fatal("expected error from panicking tool, got nil") } - if !strings.Contains(err.Error(), "boom") { - t.Fatalf("error should contain panic message, got: %s", err.Error()) + if strings.Contains(err.Error(), "boom") || strings.Contains(err.Error(), "goroutine") { + t.Fatalf("external error leaked panic details: %s", err.Error()) } - if !strings.Contains(err.Error(), "tool panic_tool panic") { - t.Fatalf("error should identify the tool, got: %s", err.Error()) + if !strings.Contains(err.Error(), "panic_tool") || !strings.Contains(err.Error(), "call-panic") { + t.Fatalf("error should identify the tool call, got: %s", err.Error()) + } + if got := logs.String(); !strings.Contains(got, "boom") || !strings.Contains(got, "goroutine") || !strings.Contains(got, "call-panic") { + t.Fatalf("panic log missing details: %s", got) } if tool.ResultText(result) != "" { t.Fatalf("result should be empty on panic, got: %s", tool.ResultText(result)) diff --git a/pkg/commands/command.go b/pkg/commands/command.go index b6551cb5..1841259f 100644 --- a/pkg/commands/command.go +++ b/pkg/commands/command.go @@ -39,14 +39,19 @@ type CommandRegistry struct { items map[string]Command order []string groups map[string][]string + logger telemetry.Logger tools map[string]tool.Tool toolOrder []string } func (r *CommandRegistry) SetLogger(logger telemetry.Logger) { + if logger == nil { + logger = telemetry.NopLogger() + } r.mu.Lock() defer r.mu.Unlock() + r.logger = logger for _, tool := range r.tools { if aware, ok := tool.(LoggerAware); ok { aware.InitLogger(logger) @@ -59,6 +64,7 @@ func NewRegistry() *CommandRegistry { items: make(map[string]Command), groups: make(map[string][]string), tools: make(map[string]tool.Tool), + logger: telemetry.NopLogger(), } } @@ -98,11 +104,40 @@ func (r *CommandRegistry) ToolDefinitions() []*tool.Definition { return defs } -func (r *CommandRegistry) ExecuteTool(ctx context.Context, name, arguments string) (result *tool.Result, err error) { +func (r *CommandRegistry) ExecuteTool(ctx context.Context, name, arguments string) (*tool.Result, error) { + return r.executeTool(ctx, name, func(t tool.Tool) (*tool.Result, error) { + return t.Execute(ctx, arguments) + }) +} + +// ExecuteBashForeground runs the transport-facing Bash path through the same +// panic boundary as ordinary tool execution. +func (r *CommandRegistry) ExecuteBashForeground(ctx context.Context, command string, options BashExecOptions) (*tool.Result, error) { + return r.executeTool(ctx, "bash", func(t tool.Tool) (*tool.Result, error) { + foreground, ok := t.(interface { + RunForegroundTool(context.Context, string, BashExecOptions) (*tool.Result, error) + }) + if !ok { + return nil, fmt.Errorf("bash tool does not support foreground execution") + } + return foreground.RunForegroundTool(ctx, command, options) + }) +} + +func (r *CommandRegistry) executeTool(ctx context.Context, name string, run func(tool.Tool) (*tool.Result, error)) (result *tool.Result, err error) { defer func() { if recovered := recover(); recovered != nil { result = nil - err = fmt.Errorf("tool %s panic: %v\n%s", name, recovered, debug.Stack()) + invocation := tool.InvocationFromContext(ctx) + r.currentLogger().Errorf( + "tool panic name=%s call_id=%s session_id=%s turn_id=%s panic=%v\n%s", + name, invocation.CallID, invocation.SessionID, invocation.TurnID, recovered, debug.Stack(), + ) + if invocation.CallID != "" { + err = fmt.Errorf("tool %s failed unexpectedly (call_id=%s)", name, invocation.CallID) + } else { + err = fmt.Errorf("tool %s failed unexpectedly", name) + } } }() @@ -110,7 +145,17 @@ func (r *CommandRegistry) ExecuteTool(ctx context.Context, name, arguments strin if !ok { return nil, fmt.Errorf("unknown tool: %s", name) } - return t.Execute(ctx, arguments) + return run(t) +} + +func (r *CommandRegistry) currentLogger() telemetry.Logger { + r.mu.RLock() + logger := r.logger + r.mu.RUnlock() + if logger == nil { + return telemetry.NopLogger() + } + return logger } func (r *CommandRegistry) Register(cmd Command, group string) { diff --git a/pkg/commands/image_optimize.go b/pkg/commands/image_optimize.go index 66e99399..048f8181 100644 --- a/pkg/commands/image_optimize.go +++ b/pkg/commands/image_optimize.go @@ -1,164 +1,33 @@ package commands import ( - "bytes" - "fmt" "image" - "image/jpeg" - "image/png" "io" - "golang.org/x/image/draw" - _ "golang.org/x/image/webp" + "github.com/chainreactors/aiscan/pkg/imageutil" ) const ( - maxDimension = 2000 - maxPayloadBytes = 3_400_000 // raw bytes; ~4.5MB base64, below Anthropic's 5MB limit + maxDimension = imageutil.MaxDimension + maxPayloadBytes = imageutil.MaxPayloadBytes ) -var jpegQualities = []int{85, 70, 55, 40} +var jpegQualities = imageutil.JPEGQualities -type optimizedImage struct { - MimeType string - Data []byte - OrigW int - OrigH int - FinalW int - FinalH int -} +type optimizedImage = imageutil.Optimized func optimizeImage(r io.Reader, srcMime string) (*optimizedImage, error) { - raw, err := io.ReadAll(r) - if err != nil { - return nil, err - } - - // GIF: pass through without decoding (may be animated) - if srcMime == "image/gif" { - return passthrough(raw, srcMime) - } - - img, _, err := image.Decode(bytes.NewReader(raw)) - if err != nil { - return passthrough(raw, srcMime) - } - - bounds := img.Bounds() - origW, origH := bounds.Dx(), bounds.Dy() - - img = resizeIfNeeded(img, origW, origH) - finalBounds := img.Bounds() - finalW, finalH := finalBounds.Dx(), finalBounds.Dy() - - data, mime, err := pickSmallestEncoding(img) - if err != nil { - return nil, err - } - - return &optimizedImage{ - MimeType: mime, - Data: data, - OrigW: origW, - OrigH: origH, - FinalW: finalW, - FinalH: finalH, - }, nil -} - -func passthrough(raw []byte, mime string) (*optimizedImage, error) { - if len(raw) > maxPayloadBytes { - return nil, fmt.Errorf("image too large after encoding (%d bytes, max %d)", len(raw), maxPayloadBytes) - } - return &optimizedImage{ - MimeType: mime, - Data: raw, - }, nil + return imageutil.Optimize(r, srcMime) } func resizeIfNeeded(img image.Image, w, h int) image.Image { - if w <= maxDimension && h <= maxDimension { - return img - } - - var newW, newH int - if w > h { - newW = maxDimension - newH = h * maxDimension / w - } else { - newH = maxDimension - newW = w * maxDimension / h - } - if newW < 1 { - newW = 1 - } - if newH < 1 { - newH = 1 - } - - dst := image.NewRGBA(image.Rect(0, 0, newW, newH)) - draw.CatmullRom.Scale(dst, dst.Bounds(), img, img.Bounds(), draw.Over, nil) - return dst -} - -// pickSmallestEncoding tries PNG and multiple JPEG quality levels, -// returning the smallest encoding that fits under maxPayloadBytes. -func pickSmallestEncoding(img image.Image) (data []byte, mime string, err error) { - pngData := encodePNG(img) - jpegData := encodeJPEG(img, jpegQualities[0]) - - // Pick smaller of PNG vs best-quality JPEG - best := pngData - bestMime := "image/png" - if len(jpegData) < len(best) { - best = jpegData - bestMime = "image/jpeg" - } - - if len(best) <= maxPayloadBytes { - return best, bestMime, nil - } - - // Too large — try lower JPEG qualities - for _, q := range jpegQualities[1:] { - jpegData = encodeJPEG(img, q) - if len(jpegData) <= maxPayloadBytes { - return jpegData, "image/jpeg", nil - } - } - - // Still too large — progressively shrink dimensions - bounds := img.Bounds() - w, h := bounds.Dx(), bounds.Dy() - for w > 1 && h > 1 { - w = w * 3 / 4 - h = h * 3 / 4 - if w < 1 { - w = 1 - } - if h < 1 { - h = 1 - } - dst := image.NewRGBA(image.Rect(0, 0, w, h)) - draw.CatmullRom.Scale(dst, dst.Bounds(), img, img.Bounds(), draw.Over, nil) - jpegData = encodeJPEG(dst, jpegQualities[0]) - if len(jpegData) <= maxPayloadBytes { - return jpegData, "image/jpeg", nil - } - } - - return nil, "", fmt.Errorf("cannot compress image to fit %d byte limit", maxPayloadBytes) + return imageutil.ResizeIfNeeded(img, w, h) } func encodePNG(img image.Image) []byte { - var buf bytes.Buffer - enc := &png.Encoder{CompressionLevel: png.BestCompression} - _ = enc.Encode(&buf, img) - return buf.Bytes() + return imageutil.EncodePNG(img) } func encodeJPEG(img image.Image, quality int) []byte { - var buf bytes.Buffer - _ = jpeg.Encode(&buf, img, &jpeg.Options{Quality: quality}) - return buf.Bytes() + return imageutil.EncodeJPEG(img, quality) } diff --git a/pkg/headless/action_types.go b/pkg/headless/action_types.go index 494c95ee..c2d458e0 100644 --- a/pkg/headless/action_types.go +++ b/pkg/headless/action_types.go @@ -41,38 +41,81 @@ const ( ActionWaitVisible // wait for element visibility ActionDialog // handle JS dialog (deprecated, use waitdialog) ActionWaitDialog // wait for JS dialog and capture type+message + + // AIScan extensions. Keep these appended so the nuclei-compatible values + // above remain stable. + ActionDblClick // double-click an element + ActionHover // hover an element + ActionFocus // focus an element + ActionBlur // blur an element + ActionCheck // ensure a checkbox/radio is checked + ActionUncheck // ensure a checkbox is unchecked + ActionDispatchEvent // dispatch a DOM event + ActionSetViewport // set viewport dimensions + ActionWaitURL // wait for the page URL to match + ActionWaitRequest // wait for a captured request URL to match + ActionWaitResponse // wait for a captured response URL to match + ActionStorage // mutate localStorage/sessionStorage + ActionCookie // mutate browser cookies + ActionAssert // assert rendered page state + ActionScroll // scroll the page mouse wheel + ActionDrag // drag one element to another + ActionReload // reload the current page + ActionGoBack // navigate backward + ActionGoForward // navigate forward + ActionSetContent // replace the current document content ) var actionTypeNames = map[ActionType]string{ - ActionNavigate: "navigate", - ActionScript: "script", - ActionClick: "click", - ActionRightClick: "rightclick", - ActionTextInput: "text", - ActionScreenshot: "screenshot", - ActionTimeInput: "time", - ActionSelectInput: "select", - ActionFilesInput: "files", - ActionWaitDOM: "waitdom", - ActionWaitFCP: "waitfcp", - ActionWaitFMP: "waitfmp", - ActionWaitIdle: "waitidle", - ActionWaitLoad: "waitload", - ActionWaitStable: "waitstable", - ActionGetResource: "getresource", - ActionExtract: "extract", - ActionSetMethod: "setmethod", - ActionAddHeader: "addheader", - ActionSetHeader: "setheader", - ActionDeleteHeader: "deleteheader", - ActionSetBody: "setbody", - ActionWaitEvent: "waitevent", - ActionKeyboard: "keyboard", - ActionDebug: "debug", - ActionSleep: "sleep", - ActionWaitVisible: "waitvisible", - ActionDialog: "dialog", - ActionWaitDialog: "waitdialog", + ActionNavigate: "navigate", + ActionScript: "script", + ActionClick: "click", + ActionRightClick: "rightclick", + ActionTextInput: "text", + ActionScreenshot: "screenshot", + ActionTimeInput: "time", + ActionSelectInput: "select", + ActionFilesInput: "files", + ActionWaitDOM: "waitdom", + ActionWaitFCP: "waitfcp", + ActionWaitFMP: "waitfmp", + ActionWaitIdle: "waitidle", + ActionWaitLoad: "waitload", + ActionWaitStable: "waitstable", + ActionGetResource: "getresource", + ActionExtract: "extract", + ActionSetMethod: "setmethod", + ActionAddHeader: "addheader", + ActionSetHeader: "setheader", + ActionDeleteHeader: "deleteheader", + ActionSetBody: "setbody", + ActionWaitEvent: "waitevent", + ActionKeyboard: "keyboard", + ActionDebug: "debug", + ActionSleep: "sleep", + ActionWaitVisible: "waitvisible", + ActionDialog: "dialog", + ActionWaitDialog: "waitdialog", + ActionDblClick: "dblclick", + ActionHover: "hover", + ActionFocus: "focus", + ActionBlur: "blur", + ActionCheck: "check", + ActionUncheck: "uncheck", + ActionDispatchEvent: "dispatch", + ActionSetViewport: "setviewport", + ActionWaitURL: "waiturl", + ActionWaitRequest: "waitrequest", + ActionWaitResponse: "waitresponse", + ActionStorage: "storage", + ActionCookie: "cookie", + ActionAssert: "assert", + ActionScroll: "scroll", + ActionDrag: "drag", + ActionReload: "reload", + ActionGoBack: "goback", + ActionGoForward: "goforward", + ActionSetContent: "setcontent", } var actionTypeMapping = func() map[string]ActionType { diff --git a/pkg/headless/engine.go b/pkg/headless/engine.go index 1ad16d41..a562f16f 100644 --- a/pkg/headless/engine.go +++ b/pkg/headless/engine.go @@ -6,10 +6,12 @@ package headless import ( + "fmt" "net/http" "sync" "time" + browserutil "github.com/chainreactors/aiscan/pkg/browser" "github.com/go-rod/rod" "github.com/go-rod/rod/lib/launcher" "github.com/go-rod/rod/lib/proto" @@ -100,6 +102,13 @@ func (e *Engine) Init() error { Set("disable-notifications"). Set("mute-audio"). Set("window-size", "1920,1080") + binary, err := browserutil.Discover() + if err != nil { + return fmt.Errorf("headless: browser discovery failed: %w", err) + } + if binary.Path != "" { + l = l.Bin(binary.Path) + } if e.options.Proxy != "" { l = l.Set("proxy-server", e.options.Proxy) diff --git a/pkg/headless/engine_test.go b/pkg/headless/engine_test.go index 94309b76..173f8d9d 100644 --- a/pkg/headless/engine_test.go +++ b/pkg/headless/engine_test.go @@ -334,6 +334,141 @@ func TestExecMultipleHeadlessRequests(t *testing.T) { } } +func TestExecAIScanExtendedActions(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("/actions", func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprint(w, ` + + + +
+
Source
+
Drop
+
+`) + }) + srv := httptest.NewServer(mux) + defer srv.Close() + + rodPage, err := sharedEngine.NewPage() + if err != nil { + t.Fatal(err) + } + page := NewPage(rodPage, sharedEngine, nil) + defer page.Close() + + action := func(kind ActionType, data map[string]string) *Action { + return &Action{ActionType: ActionTypeHolder{ActionType: kind}, Data: data} + } + actions := []*Action{ + action(ActionNavigate, map[string]string{"url": srv.URL + "/actions"}), + action(ActionWaitURL, map[string]string{"url": "/actions"}), + action(ActionWaitRequest, map[string]string{"url": "/actions"}), + action(ActionWaitResponse, map[string]string{"url": "/actions"}), + action(ActionTextInput, mergeMapsForTest(ParseSelector("label=Email"), map[string]string{"value": "alice@example.com", "clear": "true"})), + action(ActionKeyboard, mergeMapsForTest(ParseSelector("label=Email"), map[string]string{"keys": "End"})), + action(ActionFocus, ParseSelector("label=Email")), + action(ActionBlur, ParseSelector("label=Email")), + action(ActionCheck, ParseSelector("testid=terms")), + action(ActionCheck, ParseSelector("testid=terms")), + action(ActionUncheck, ParseSelector("testid=terms")), + action(ActionCheck, ParseSelector("testid=terms")), + action(ActionSelectInput, mergeMapsForTest(ParseSelector(`role=combobox[name="Plan"]`), map[string]string{"value": "pro"})), + action(ActionHover, ParseSelector(`role=button[name="Activate"]`)), + action(ActionDblClick, ParseSelector(`role=button[name="Activate"]`)), + action(ActionDispatchEvent, mergeMapsForTest(ParseSelector("#activate"), map[string]string{"event": "aiscan", "detail": `{"flag":"ok"}`})), + action(ActionDrag, mergeMapsForTest(ParseSelector("testid=source"), map[string]string{"target": "testid=drop"})), + action(ActionScroll, map[string]string{"y": "250", "steps": "2"}), + action(ActionStorage, map[string]string{"storage": "local", "operation": "set", "key": "token", "value": "abc123"}), + action(ActionCookie, map[string]string{"operation": "set", "name": "session", "value": "cookie-value"}), + action(ActionSetViewport, map[string]string{"width": "1024", "height": "768"}), + action(ActionAssert, mergeMapsForTest(ParseSelector("label=Email"), map[string]string{"type": "value", "value": "alice@example.com"})), + action(ActionAssert, mergeMapsForTest(ParseSelector("testid=terms"), map[string]string{"type": "checked"})), + action(ActionAssert, mergeMapsForTest(ParseSelector(`role=combobox[name="Plan"]`), map[string]string{"type": "value", "value": "pro"})), + action(ActionAssert, mergeMapsForTest(ParseSelector("testid=state"), map[string]string{"type": "attribute", "attribute": "data-focus", "value": "yes"})), + action(ActionAssert, mergeMapsForTest(ParseSelector("testid=state"), map[string]string{"type": "attribute", "attribute": "data-blur", "value": "yes"})), + action(ActionAssert, mergeMapsForTest(ParseSelector("testid=state"), map[string]string{"type": "attribute", "attribute": "data-hover", "value": "yes"})), + action(ActionAssert, mergeMapsForTest(ParseSelector("testid=state"), map[string]string{"type": "attribute", "attribute": "data-dblclick", "value": "yes"})), + action(ActionAssert, mergeMapsForTest(ParseSelector("testid=state"), map[string]string{"type": "attribute", "attribute": "data-custom", "value": "ok"})), + action(ActionAssert, mergeMapsForTest(ParseSelector("testid=state"), map[string]string{"type": "attribute", "attribute": "data-dragstart", "value": "yes"})), + action(ActionAssert, mergeMapsForTest(ParseSelector("testid=state"), map[string]string{"type": "attribute", "attribute": "data-dragend", "value": "yes"})), + action(ActionAssert, map[string]string{"type": "storage", "storage": "local", "key": "token", "value": "abc123"}), + action(ActionAssert, map[string]string{"type": "cookie", "name": "session", "value": "cookie-value"}), + action(ActionSetContent, map[string]string{"html": `
Replacement content
`}), + action(ActionAssert, mergeMapsForTest(ParseSelector("testid=replacement"), map[string]string{"type": "text", "value": "Replacement content"})), + } + if _, err := page.ExecuteActions(actions); err != nil { + t.Fatalf("extended action replay failed: %v", err) + } + + viewport, err := rodPage.Eval(`() => [window.innerWidth, window.innerHeight]`) + if err != nil { + t.Fatal(err) + } + if got := viewport.Value.Arr(); len(got) != 2 || got[0].Int() != 1024 || got[1].Int() != 768 { + t.Fatalf("viewport = %v, want 1024x768", viewport.Value.Val()) + } +} + +func TestExecAIScanHistoryActions(t *testing.T) { + mux := http.NewServeMux() + mux.HandleFunc("/one", func(w http.ResponseWriter, _ *http.Request) { + fmt.Fprint(w, `Page Oneone`) + }) + mux.HandleFunc("/two", func(w http.ResponseWriter, _ *http.Request) { + fmt.Fprint(w, `Page Twotwo`) + }) + srv := httptest.NewServer(mux) + defer srv.Close() + + rodPage, err := sharedEngine.NewPage() + if err != nil { + t.Fatal(err) + } + page := NewPage(rodPage, sharedEngine, nil) + defer page.Close() + action := func(kind ActionType, data map[string]string) *Action { + return &Action{ActionType: ActionTypeHolder{ActionType: kind}, Data: data} + } + + actions := []*Action{ + action(ActionNavigate, map[string]string{"url": srv.URL + "/one"}), + action(ActionNavigate, map[string]string{"url": srv.URL + "/two"}), + action(ActionGoBack, map[string]string{}), + action(ActionAssert, map[string]string{"type": "url", "value": "/one", "match": "contains"}), + action(ActionGoForward, map[string]string{}), + action(ActionAssert, map[string]string{"type": "url", "value": "/two", "match": "contains"}), + action(ActionReload, map[string]string{}), + action(ActionAssert, map[string]string{"type": "title", "value": "Page Two"}), + } + if _, err := page.ExecuteActions(actions); err != nil { + t.Fatalf("history action replay failed: %v", err) + } +} + +func mergeMapsForTest(left, right map[string]string) map[string]string { + merged := make(map[string]string, len(left)+len(right)) + for key, value := range left { + merged[key] = value + } + for key, value := range right { + merged[key] = value + } + return merged +} + // ========================================================================== // Engine lifecycle // ========================================================================== diff --git a/pkg/headless/page.go b/pkg/headless/page.go index aadda6eb..620395b9 100644 --- a/pkg/headless/page.go +++ b/pkg/headless/page.go @@ -27,6 +27,9 @@ const ( type HistoryEntry struct { RawRequest string RawResponse string + URL string + Method string + StatusCode int } // Page wraps a go-rod page and executes headless action sequences. @@ -157,7 +160,7 @@ func (p *Page) ExecuteActions(actions []*Action) (ActionData, error) { case ActionWaitFMP: err = p.actionWaitLifecycle(resolved, out, proto.PageLifecycleEventNameFirstMeaningfulPaint) case ActionWaitIdle: - err = p.actionWaitLifecycle(resolved, out, proto.PageLifecycleEventNameNetworkIdle) + err = p.actionWaitIdle(resolved, out) case ActionWaitLoad: err = p.actionWaitLifecycle(resolved, out, proto.PageLifecycleEventNameLoad) case ActionWaitStable: @@ -186,6 +189,49 @@ func (p *Page) ExecuteActions(actions []*Action) (ActionData, error) { err = p.actionDialog(resolved, out) case ActionWaitDialog: err = p.actionWaitDialog(resolved, out) + case ActionDblClick: + err = p.actionDblClick(resolved, out) + case ActionHover: + err = p.actionHover(resolved, out) + case ActionFocus: + err = p.actionFocus(resolved, out) + case ActionBlur: + err = p.actionBlur(resolved, out) + case ActionCheck: + err = p.actionCheck(resolved, out, true) + case ActionUncheck: + err = p.actionCheck(resolved, out, false) + case ActionDispatchEvent: + err = p.actionDispatchEvent(resolved, out) + case ActionSetViewport: + err = p.actionSetViewport(resolved, out) + case ActionWaitURL: + err = p.actionWaitURL(resolved, out) + case ActionWaitRequest: + err = p.actionWaitNetwork(resolved, out, false) + case ActionWaitResponse: + err = p.actionWaitNetwork(resolved, out, true) + case ActionStorage: + err = p.actionStorage(resolved, out) + case ActionCookie: + err = p.actionCookie(resolved, out) + case ActionAssert: + err = p.actionAssert(resolved, out) + if err != nil { + err = fmt.Errorf("%s assertion: %w", firstNonEmpty(resolved.GetArg("type"), resolved.GetArg("target")), err) + } + case ActionScroll: + err = p.actionScroll(resolved, out) + case ActionDrag: + err = p.actionDrag(resolved, out) + case ActionReload: + err = p.actionReload(resolved, out) + case ActionGoBack: + err = p.actionHistoryNavigation(resolved, out, false) + case ActionGoForward: + err = p.actionHistoryNavigation(resolved, out, true) + case ActionSetContent: + err = p.actionSetContent(resolved, out) default: continue } @@ -350,7 +396,7 @@ func (p *Page) captureHijackHistory(ctx *rod.Hijack) { rawResp.WriteString(ctx.Response.Body()) } - p.addHistory(rawReq, rawResp.String(), payload) + p.addHistory(rawReq, rawResp.String(), req.Method, req.URL.String(), payload) } // routingRuleHandlerNative handles capture-only interception via native CDP Fetch. @@ -393,6 +439,9 @@ func (p *Page) routingRuleHandlerNative(e *proto.FetchRequestPaused) error { p.History = append(p.History, HistoryEntry{ RawRequest: rawReq.String(), RawResponse: rawResp.String(), + URL: e.Request.URL, + Method: e.Request.Method, + StatusCode: statusCode, }) p.mu.Unlock() @@ -400,7 +449,7 @@ func (p *Page) routingRuleHandlerNative(e *proto.FetchRequestPaused) error { } // addHistory records a request/response pair from the HijackRouter path. -func (p *Page) addHistory(rawReq, rawResp string, payload *proto.FetchFulfillRequest) { +func (p *Page) addHistory(rawReq, rawResp, method, requestURL string, payload *proto.FetchFulfillRequest) { p.mu.Lock() defer p.mu.Unlock() @@ -410,10 +459,16 @@ func (p *Page) addHistory(rawReq, rawResp string, payload *proto.FetchFulfillReq p.responseHeaders[h.Name] = h.Value } } - p.History = append(p.History, HistoryEntry{ + entry := HistoryEntry{ RawRequest: rawReq, RawResponse: rawResp, - }) + URL: requestURL, + Method: method, + } + if payload != nil { + entry.StatusCode = payload.ResponseCode + } + p.History = append(p.History, entry) } // Close cleans up any resources held by the page. @@ -430,35 +485,7 @@ func (p *Page) Close() { // pageElementBy resolves a page element using nuclei's selector conventions. func (p *Page) pageElementBy(data map[string]string) (*rod.Element, error) { - by := data["by"] - page := p.page.Timeout(defaultActionTimeout) - switch by { - case "x", "xpath": - xpath := data["xpath"] - if xpath == "" { - return nil, fmt.Errorf("xpath selector required") - } - return page.ElementX(xpath) - case "js": - return page.ElementByJS(&rod.EvalOptions{JS: data["js"]}) - case "r", "regex": - return page.ElementR(data["selector"], data["regex"]) - case "search": - elms, err := page.Search(data["query"]) - if err != nil { - return nil, err - } - if elms.First != nil { - return elms.First, nil - } - return nil, fmt.Errorf("no element found for query: %s", data["query"]) - default: - sel := data["selector"] - if sel == "" { - return nil, fmt.Errorf("no selector provided") - } - return page.Element(sel) - } + return ElementBy(p.page, data, defaultActionTimeout) } // ResponseData captures HTTP response info from the navigated page. diff --git a/pkg/headless/page_actions.go b/pkg/headless/page_actions.go index 0ed11777..df33a292 100644 --- a/pkg/headless/page_actions.go +++ b/pkg/headless/page_actions.go @@ -17,7 +17,6 @@ import ( "time" "github.com/go-rod/rod" - "github.com/go-rod/rod/lib/input" "github.com/go-rod/rod/lib/proto" ) @@ -110,8 +109,8 @@ func (p *Page) actionRightClick(act *Action, out ActionData) error { } func (p *Page) actionTextInput(act *Action, out ActionData) error { - value := act.GetArg("value") - if value == "" { + value, ok := act.Data["value"] + if !ok { return fmt.Errorf("text: value argument required") } el, err := p.pageElementBy(act.Data) @@ -121,6 +120,11 @@ func (p *Page) actionTextInput(act *Action, out ActionData) error { if err := el.ScrollIntoView(); err != nil { return fmt.Errorf("text scroll: %w", err) } + if act.GetArg("clear") == "true" { + if err := el.SelectAllText(); err != nil { + return fmt.Errorf("text clear: %w", err) + } + } return el.Input(value) } @@ -130,8 +134,18 @@ func (p *Page) actionScreenshot(act *Action, out ActionData) error { to = "screenshot" } - fullpage := act.GetArg("fullpage") == "true" - data, err := p.page.Screenshot(fullpage, &proto.PageCaptureScreenshot{}) + var data []byte + var err error + if hasSelectorArgs(act.Data) { + var el *rod.Element + el, err = p.pageElementBy(act.Data) + if err == nil { + data, err = el.Screenshot(proto.PageCaptureScreenshotFormatPng, 90) + } + } else { + fullpage := act.GetArg("fullpage") == "true" + data, err = p.page.Screenshot(fullpage, &proto.PageCaptureScreenshot{}) + } if err != nil { return fmt.Errorf("screenshot: %w", err) } @@ -187,9 +201,12 @@ func (p *Page) actionSelectInput(act *Action, out ActionData) error { if err := el.ScrollIntoView(); err != nil { return fmt.Errorf("select scroll: %w", err) } - selected := act.GetArg("selected") == "true" + selected := !strings.EqualFold(act.GetArg("selected"), "false") selectorType := selectorBy(act.GetArg("selector")) - return el.Select([]string{value}, selected, selectorType) + if err := el.Select([]string{value}, selected, selectorType); err == nil { + return nil + } + return el.Select([]string{fmt.Sprintf("option[value=%s]", strconv.Quote(value))}, selected, rod.SelectorTypeCSSSector) } func (p *Page) actionFilesInput(act *Action, out ActionData) error { @@ -226,13 +243,24 @@ func (p *Page) actionWaitStable(act *Action, out ActionData) error { return p.page.Timeout(timeout).WaitStable(dur) } +func (p *Page) actionWaitIdle(act *Action, out ActionData) error { + idle := 500 * time.Millisecond + if value := act.GetArg("duration"); value != "" { + if parsed, err := time.ParseDuration(value); err == nil { + idle = parsed + } + } + wait := p.page.Timeout(p.getTimeout(act)).WaitRequestIdle(idle, nil, nil, nil) + wait() + return nil +} + func (p *Page) actionWaitVisible(act *Action, out ActionData) error { - sel := act.GetArg("selector") - if sel == "" { + if !hasSelectorArgs(act.Data) { return fmt.Errorf("waitvisible: selector argument required") } timeout := p.getTimeout(act) - el, err := p.page.Timeout(timeout).Element(sel) + el, err := ElementBy(p.page, act.Data, timeout) if err != nil { return fmt.Errorf("waitvisible: %w", err) } @@ -255,6 +283,42 @@ func (p *Page) actionGetResource(act *Action, out ActionData) error { } func (p *Page) actionExtract(act *Action, out ActionData) error { + target := act.GetArg("target") + if target == "url" || target == "title" { + info, err := p.page.Info() + if err != nil { + return fmt.Errorf("extract %s: %w", target, err) + } + value := info.URL + if target == "title" { + value = info.Title + } + if act.Name != "" { + out[act.Name] = value + } + return nil + } + if target == "storage" { + value, err := p.readStorage(act.GetArg("storage"), act.GetArg("key")) + if err != nil { + return err + } + if act.Name != "" { + out[act.Name] = value + } + return nil + } + if target == "cookie" { + value, err := p.readCookie(act.GetArg("name")) + if err != nil { + return err + } + if act.Name != "" { + out[act.Name] = value + } + return nil + } + el, err := p.pageElementBy(act.Data) if err != nil { return fmt.Errorf("extract: %w", err) @@ -263,7 +327,6 @@ func (p *Page) actionExtract(act *Action, out ActionData) error { return fmt.Errorf("extract scroll: %w", err) } - target := act.GetArg("target") switch target { case "attribute": attr := act.GetArg("attribute") @@ -281,6 +344,34 @@ func (p *Page) actionExtract(act *Action, out ActionData) error { out[act.Name] = "" } } + case "html": + html, err := el.HTML() + if err != nil { + return err + } + if act.Name != "" { + out[act.Name] = html + } + case "value": + value, err := el.Property("value") + if err != nil { + return err + } + if act.Name != "" { + out[act.Name] = value.String() + } + case "property": + property := act.GetArg("property") + if property == "" { + return fmt.Errorf("extract: property name required") + } + value, err := el.Property(property) + if err != nil { + return err + } + if act.Name != "" { + out[act.Name] = value.Val() + } default: text, err := el.Text() if err != nil { @@ -298,7 +389,16 @@ func (p *Page) actionKeyboard(act *Action, out ActionData) error { if keys == "" { return fmt.Errorf("keyboard: keys argument required") } - return p.page.Keyboard.Type([]input.Key(keys)...) + if hasSelectorArgs(act.Data) { + el, err := p.pageElementBy(act.Data) + if err != nil { + return fmt.Errorf("keyboard selector: %w", err) + } + if err := el.Focus(); err != nil { + return fmt.Errorf("keyboard focus: %w", err) + } + } + return pressKeys(p.page, keys) } func (p *Page) actionSleep(act *Action, out ActionData) error { @@ -369,9 +469,11 @@ func (p *Page) actionWaitEvent(act *Action, out ActionData) (func() error, error func (p *Page) actionDialog(act *Action, out ActionData) error { wait, handle := p.page.MustHandleDialog() + accept := !strings.EqualFold(act.GetArg("accept"), "false") + prompt := act.GetArg("prompt") go func() { wait() - handle(true, "") + handle(accept, prompt) }() return nil } @@ -395,11 +497,13 @@ func (p *Page) actionWaitDialog(act *Action, out ActionData) error { ch := make(chan dialogResult, 1) wait, handle := p.page.HandleDialog() + accept := !strings.EqualFold(act.GetArg("accept"), "false") + prompt := act.GetArg("prompt") go func() { dialog := wait() err := handle(&proto.PageHandleJavaScriptDialog{ - Accept: true, - PromptText: "", + Accept: accept, + PromptText: prompt, }) ch <- dialogResult{dialog: dialog, err: err} }() diff --git a/pkg/headless/page_actions_extended.go b/pkg/headless/page_actions_extended.go new file mode 100644 index 00000000..378a212f --- /dev/null +++ b/pkg/headless/page_actions_extended.go @@ -0,0 +1,628 @@ +//go:build full + +package headless + +import ( + "encoding/json" + "fmt" + "regexp" + "strconv" + "strings" + "time" + + "github.com/go-rod/rod" + "github.com/go-rod/rod/lib/input" + "github.com/go-rod/rod/lib/proto" +) + +var headlessKeyNames = map[string]input.Key{ + "enter": input.Enter, "tab": input.Tab, "escape": input.Escape, + "backspace": input.Backspace, "delete": input.Delete, "space": input.Space, + "arrowup": input.ArrowUp, "arrowdown": input.ArrowDown, + "arrowleft": input.ArrowLeft, "arrowright": input.ArrowRight, + "home": input.Home, "end": input.End, + "pageup": input.PageUp, "pagedown": input.PageDown, + "insert": input.Insert, + "f1": input.F1, "f2": input.F2, "f3": input.F3, "f4": input.F4, + "f5": input.F5, "f6": input.F6, "f7": input.F7, "f8": input.F8, + "f9": input.F9, "f10": input.F10, "f11": input.F11, "f12": input.F12, + "shift": input.ShiftLeft, "control": input.ControlLeft, "ctrl": input.ControlLeft, + "alt": input.AltLeft, "meta": input.MetaLeft, "command": input.MetaLeft, +} + +func hasSelectorArgs(data map[string]string) bool { + if data == nil { + return false + } + return data["selector"] != "" || data["xpath"] != "" || data["js"] != "" || + data["query"] != "" || data["role"] != "" || data["label"] != "" || + data["text"] != "" || data["testid"] != "" +} + +func resolveHeadlessKey(name string) (input.Key, error) { + name = strings.TrimSpace(name) + if key, ok := headlessKeyNames[strings.ToLower(name)]; ok { + return key, nil + } + runes := []rune(name) + if len(runes) == 1 { + return input.Key(runes[0]), nil + } + return 0, fmt.Errorf("unknown key %q", name) +} + +func pressKeys(page *rod.Page, expression string) error { + parts := strings.Split(expression, "+") + if len(parts) == 1 { + key, err := resolveHeadlessKey(parts[0]) + if err != nil { + return err + } + return page.Keyboard.Type(key) + } + + actions := page.KeyActions() + modifiers := make([]input.Key, 0, len(parts)-1) + for _, part := range parts[:len(parts)-1] { + key, err := resolveHeadlessKey(part) + if err != nil { + return fmt.Errorf("modifier: %w", err) + } + modifiers = append(modifiers, key) + actions = actions.Press(key) + } + main, err := resolveHeadlessKey(parts[len(parts)-1]) + if err != nil { + return err + } + actions = actions.Type(main) + for i := len(modifiers) - 1; i >= 0; i-- { + actions = actions.Release(modifiers[i]) + } + return actions.Do() +} + +func (p *Page) actionDblClick(act *Action, _ ActionData) error { + el, err := p.pageElementBy(act.Data) + if err != nil { + return fmt.Errorf("dblclick: %w", err) + } + return el.Click(proto.InputMouseButtonLeft, 2) +} + +func (p *Page) actionHover(act *Action, _ ActionData) error { + el, err := p.pageElementBy(act.Data) + if err != nil { + return fmt.Errorf("hover: %w", err) + } + return el.Hover() +} + +func (p *Page) actionFocus(act *Action, _ ActionData) error { + el, err := p.pageElementBy(act.Data) + if err != nil { + return fmt.Errorf("focus: %w", err) + } + return el.Focus() +} + +func (p *Page) actionBlur(act *Action, _ ActionData) error { + el, err := p.pageElementBy(act.Data) + if err != nil { + return fmt.Errorf("blur: %w", err) + } + return el.Blur() +} + +func (p *Page) actionCheck(act *Action, _ ActionData, checked bool) error { + el, err := p.pageElementBy(act.Data) + if err != nil { + return fmt.Errorf("checkbox: %w", err) + } + current, err := el.Property("checked") + if err != nil { + return fmt.Errorf("checkbox state: %w", err) + } + if current.Bool() == checked { + return nil + } + if err := el.Click(proto.InputMouseButtonLeft, 1); err != nil { + return fmt.Errorf("checkbox click: %w", err) + } + current, err = el.Property("checked") + if err != nil { + return fmt.Errorf("checkbox verify: %w", err) + } + if current.Bool() != checked { + return fmt.Errorf("checkbox did not become checked=%t", checked) + } + return nil +} + +func (p *Page) actionDispatchEvent(act *Action, _ ActionData) error { + eventType := act.GetArg("event") + if eventType == "" { + eventType = act.GetArg("type") + } + if eventType == "" { + return fmt.Errorf("dispatch: event argument required") + } + el, err := p.pageElementBy(act.Data) + if err != nil { + return fmt.Errorf("dispatch: %w", err) + } + detail := act.GetArg("detail") + if detail == "" { + detail = "null" + } else if !json.Valid([]byte(detail)) { + encoded, _ := json.Marshal(detail) + detail = string(encoded) + } + _, err = el.Eval(`(eventType, detailJSON) => { + const detail = JSON.parse(detailJSON); + const options = {bubbles: true, cancelable: true}; + const event = detail === null ? new Event(eventType, options) : new CustomEvent(eventType, {...options, detail}); + this.dispatchEvent(event); + }`, eventType, detail) + return err +} + +func (p *Page) actionSetViewport(act *Action, _ ActionData) error { + width, err := positiveInt(act.GetArg("width"), "width") + if err != nil { + return fmt.Errorf("setviewport: %w", err) + } + height, err := positiveInt(act.GetArg("height"), "height") + if err != nil { + return fmt.Errorf("setviewport: %w", err) + } + scale := 1.0 + if value := act.GetArg("device-scale-factor"); value != "" { + scale, err = strconv.ParseFloat(value, 64) + if err != nil || scale <= 0 { + return fmt.Errorf("setviewport: device-scale-factor must be positive") + } + } + return p.page.SetViewport(&proto.EmulationSetDeviceMetricsOverride{ + Width: width, Height: height, DeviceScaleFactor: scale, + }) +} + +func positiveInt(value, name string) (int, error) { + n, err := strconv.Atoi(value) + if err != nil || n <= 0 { + return 0, fmt.Errorf("%s must be a positive integer", name) + } + return n, nil +} + +func (p *Page) actionWaitURL(act *Action, _ ActionData) error { + expected := firstNonEmpty(act.GetArg("url"), act.GetArg("value")) + if expected == "" { + return fmt.Errorf("waiturl: url argument required") + } + return pollUntil(p.getTimeout(act), func() (bool, error) { + info, err := p.page.Info() + if err != nil { + return false, err + } + return matchString(info.URL, expected, act.GetArg("match")) + }) +} + +func (p *Page) actionWaitNetwork(act *Action, _ ActionData, response bool) error { + expected := firstNonEmpty(act.GetArg("url"), act.GetArg("value")) + if expected == "" { + return fmt.Errorf("network wait: url argument required") + } + method := strings.ToUpper(act.GetArg("method")) + return pollUntil(p.getTimeout(act), func() (bool, error) { + p.mu.RLock() + entries := append([]HistoryEntry(nil), p.History...) + p.mu.RUnlock() + for _, entry := range entries { + if response && entry.StatusCode == 0 { + continue + } + if method != "" && strings.ToUpper(entry.Method) != method { + continue + } + matched, err := matchString(entry.URL, expected, act.GetArg("match")) + if err != nil { + return false, err + } + if matched { + return true, nil + } + } + return false, nil + }) +} + +func pollUntil(timeout time.Duration, condition func() (bool, error)) error { + deadline := time.Now().Add(timeout) + for { + ok, err := condition() + if err != nil { + return err + } + if ok { + return nil + } + if time.Now().After(deadline) { + return fmt.Errorf("condition not met within %s", timeout) + } + time.Sleep(50 * time.Millisecond) + } +} + +func matchString(actual, expected, mode string) (bool, error) { + switch strings.ToLower(strings.TrimSpace(mode)) { + case "", "contains": + return strings.Contains(actual, expected), nil + case "equals", "equal", "exact": + return actual == expected, nil + case "regex", "regexp": + return regexp.MatchString(expected, actual) + default: + return false, fmt.Errorf("unknown match mode %q", mode) + } +} + +func normalizeStorageKind(kind string) (string, error) { + switch strings.ToLower(strings.TrimSpace(kind)) { + case "", "local", "localstorage": + return "localStorage", nil + case "session", "sessionstorage": + return "sessionStorage", nil + default: + return "", fmt.Errorf("storage type must be local or session") + } +} + +func (p *Page) actionStorage(act *Action, _ ActionData) error { + kind, err := normalizeStorageKind(firstNonEmpty(act.GetArg("storage"), act.GetArg("type"))) + if err != nil { + return err + } + operation := strings.ToLower(firstNonEmpty(act.GetArg("operation"), act.GetArg("op"), "set")) + key := act.GetArg("key") + switch operation { + case "set": + value, ok := act.Data["value"] + if key == "" || !ok { + return fmt.Errorf("storage set requires key and value") + } + _, err = p.page.Eval(`(kind, key, value) => window[kind].setItem(key, value)`, kind, key, value) + case "delete", "remove": + if key == "" { + return fmt.Errorf("storage delete requires key") + } + _, err = p.page.Eval(`(kind, key) => window[kind].removeItem(key)`, kind, key) + case "clear": + _, err = p.page.Eval(`kind => window[kind].clear()`, kind) + default: + return fmt.Errorf("unknown storage operation %q", operation) + } + return err +} + +func (p *Page) readStorage(kind, key string) (interface{}, error) { + normalized, err := normalizeStorageKind(kind) + if err != nil { + return nil, err + } + if key != "" { + result, evalErr := p.page.Eval(`(kind, key) => window[kind].getItem(key)`, normalized, key) + if evalErr != nil { + return nil, evalErr + } + if result.Value.Nil() { + return "", nil + } + return result.Value.String(), nil + } + result, err := p.page.Eval(`kind => { + const output = {}; + for (let i = 0; i < window[kind].length; i++) { + const key = window[kind].key(i); + output[key] = window[kind].getItem(key); + } + return output; + }`, normalized) + if err != nil { + return nil, err + } + return result.Value.Val(), nil +} + +func (p *Page) actionCookie(act *Action, _ ActionData) error { + operation := strings.ToLower(firstNonEmpty(act.GetArg("operation"), act.GetArg("op"), "set")) + name := act.GetArg("name") + switch operation { + case "set": + value, ok := act.Data["value"] + if name == "" || !ok { + return fmt.Errorf("cookie set requires name and value") + } + cookieURL := act.GetArg("url") + if cookieURL == "" { + info, err := p.page.Info() + if err != nil { + return err + } + cookieURL = info.URL + } + cookie := &proto.NetworkCookieParam{ + Name: name, Value: value, URL: cookieURL, + Domain: act.GetArg("domain"), Path: act.GetArg("path"), + Secure: strings.EqualFold(act.GetArg("secure"), "true"), + HTTPOnly: strings.EqualFold(act.GetArg("http-only"), "true"), + } + return p.page.SetCookies([]*proto.NetworkCookieParam{cookie}) + case "delete", "remove": + if name == "" { + return fmt.Errorf("cookie delete requires name") + } + cookies, err := p.page.Cookies(nil) + if err != nil { + return err + } + for _, cookie := range cookies { + if cookie.Name == name { + if err := (proto.NetworkDeleteCookies{Name: cookie.Name, Domain: cookie.Domain, Path: cookie.Path}).Call(p.page); err != nil { + return err + } + } + } + return nil + case "clear": + cookies, err := p.page.Cookies(nil) + if err != nil { + return err + } + for _, cookie := range cookies { + if err := (proto.NetworkDeleteCookies{Name: cookie.Name, Domain: cookie.Domain, Path: cookie.Path}).Call(p.page); err != nil { + return err + } + } + return nil + default: + return fmt.Errorf("unknown cookie operation %q", operation) + } +} + +func (p *Page) readCookie(name string) (interface{}, error) { + cookies, err := p.page.Cookies(nil) + if err != nil { + return nil, err + } + if name == "" { + values := make(map[string]string, len(cookies)) + for _, cookie := range cookies { + values[cookie.Name] = cookie.Value + } + return values, nil + } + for _, cookie := range cookies { + if cookie.Name == name { + return cookie.Value, nil + } + } + return "", nil +} + +func (p *Page) actionAssert(act *Action, _ ActionData) error { + kind := strings.ToLower(firstNonEmpty(act.GetArg("type"), act.GetArg("target"))) + expected := act.GetArg("value") + var actual interface{} + + switch kind { + case "url", "title": + info, err := p.page.Info() + if err != nil { + return err + } + actual = info.URL + if kind == "title" { + actual = info.Title + } + case "storage": + value, err := p.readStorage(act.GetArg("storage"), act.GetArg("key")) + if err != nil { + return err + } + actual = value + case "cookie": + value, err := p.readCookie(act.GetArg("name")) + if err != nil { + return err + } + actual = value + default: + if !hasSelectorArgs(act.Data) { + return fmt.Errorf("assert %s requires a selector", kind) + } + el, err := p.pageElementBy(act.Data) + if err != nil { + if kind == "hidden" { + return nil + } + return fmt.Errorf("resolve %s: %w", selectorSummary(act.Data), err) + } + switch kind { + case "visible", "hidden": + visible, err := el.Visible() + if err != nil { + return err + } + want := kind == "visible" + if visible != want { + return fmt.Errorf("expected element visible=%t", want) + } + return nil + case "checked", "unchecked": + checked, err := el.Property("checked") + if err != nil { + return err + } + want := kind == "checked" + if checked.Bool() != want { + return fmt.Errorf("expected element checked=%t", want) + } + return nil + case "enabled", "disabled": + disabled, err := el.Disabled() + if err != nil { + return err + } + wantDisabled := kind == "disabled" + if disabled != wantDisabled { + return fmt.Errorf("expected element disabled=%t", wantDisabled) + } + return nil + case "text", "": + actual, err = el.Text() + if err != nil { + return err + } + case "value": + value, err := el.Property("value") + if err != nil { + return err + } + actual = value.String() + case "attribute": + attribute := act.GetArg("attribute") + if attribute == "" { + return fmt.Errorf("assert attribute requires attribute name") + } + value, err := el.Attribute(attribute) + if err != nil { + return fmt.Errorf("read attribute %q: %w", attribute, err) + } + if value != nil { + actual = *value + } else { + actual = "" + } + default: + return fmt.Errorf("unknown assertion type %q", kind) + } + } + + actualText := fmt.Sprint(actual) + matched, err := matchString(actualText, expected, firstNonEmpty(act.GetArg("match"), "equals")) + if err != nil { + return err + } + if !matched { + return fmt.Errorf("assertion failed: got %q, expected %s %q", actualText, firstNonEmpty(act.GetArg("match"), "equals"), expected) + } + return nil +} + +func (p *Page) actionScroll(act *Action, _ ActionData) error { + x, err := parseFloatDefault(act.GetArg("x"), 0) + if err != nil { + return fmt.Errorf("scroll x: %w", err) + } + y, err := parseFloatDefault(firstNonEmpty(act.GetArg("y"), act.GetArg("delta-y")), 0) + if err != nil { + return fmt.Errorf("scroll y: %w", err) + } + steps := 1 + if raw := act.GetArg("steps"); raw != "" { + steps, err = positiveInt(raw, "steps") + if err != nil { + return err + } + } + return p.page.Mouse.Scroll(x, y, steps) +} + +func parseFloatDefault(value string, fallback float64) (float64, error) { + if value == "" { + return fallback, nil + } + return strconv.ParseFloat(value, 64) +} + +func (p *Page) actionDrag(act *Action, _ ActionData) error { + source, err := p.pageElementBy(act.Data) + if err != nil { + return fmt.Errorf("drag source: %w", err) + } + targetSelector := act.GetArg("target") + if targetSelector == "" { + return fmt.Errorf("drag target selector required") + } + target, err := FindElement(p.page, targetSelector, p.getTimeout(act)) + if err != nil { + return fmt.Errorf("drag target: %w", err) + } + if err := source.Hover(); err != nil { + return err + } + if err := p.page.Mouse.Down(proto.InputMouseButtonLeft, 1); err != nil { + return err + } + defer func() { _ = p.page.Mouse.Up(proto.InputMouseButtonLeft, 1) }() + if err := target.Hover(); err != nil { + return err + } + return p.page.Mouse.Up(proto.InputMouseButtonLeft, 1) +} + +func (p *Page) actionReload(act *Action, _ ActionData) error { + if err := p.page.Timeout(p.getTimeout(act)).Reload(); err != nil { + return err + } + return p.page.Timeout(p.getTimeout(act)).WaitStable(defaultStableDur) +} + +func (p *Page) actionHistoryNavigation(act *Action, _ ActionData, forward bool) error { + page := p.page.Timeout(p.getTimeout(act)) + var err error + if forward { + err = page.NavigateForward() + } else { + err = page.NavigateBack() + } + if err != nil { + return err + } + return page.WaitStable(defaultStableDur) +} + +func (p *Page) actionSetContent(act *Action, _ ActionData) error { + html, ok := act.Data["html"] + if !ok { + html, ok = act.Data["value"] + } + if !ok { + return fmt.Errorf("setcontent: html argument required") + } + return p.page.SetDocumentContent(html) +} + +func firstNonEmpty(values ...string) string { + for _, value := range values { + if value != "" { + return value + } + } + return "" +} + +func selectorSummary(data map[string]string) string { + by := strings.ToLower(data["by"]) + if by == "" { + return fmt.Sprintf("selector %q", data["selector"]) + } + value := data[by] + if by == "role" { + value = data["role"] + " name=" + data["name"] + } + return fmt.Sprintf("%s selector %q", by, value) +} diff --git a/pkg/headless/selector.go b/pkg/headless/selector.go new file mode 100644 index 00000000..94d478b8 --- /dev/null +++ b/pkg/headless/selector.go @@ -0,0 +1,205 @@ +//go:build full + +package headless + +import ( + "fmt" + "strings" + "time" + + "github.com/go-rod/rod" +) + +// semanticSelectorJS resolves the small, stable locator vocabulary shared by +// the CLI recorder and the headless replay engine. It traverses open shadow +// roots and implements the useful subset of Playwright-style locators without +// coupling templates to Playwright's private selector engine. +const semanticSelectorJS = `(kind, role, name, value, exact, testIdAttribute) => { + const normalize = text => String(text || '').replace(/\s+/g, ' ').trim(); + const matches = (actual, expected) => { + actual = normalize(actual); + expected = normalize(expected); + return exact ? actual === expected : actual.toLowerCase().includes(expected.toLowerCase()); + }; + const elements = []; + const visit = root => { + for (const element of root.querySelectorAll('*')) { + elements.push(element); + if (element.shadowRoot) visit(element.shadowRoot); + } + }; + visit(document); + + const implicitRole = element => { + const explicit = element.getAttribute('role'); + if (explicit) return explicit.split(/\s+/)[0].toLowerCase(); + const tag = element.tagName.toLowerCase(); + const type = (element.getAttribute('type') || '').toLowerCase(); + if (tag === 'a' && element.hasAttribute('href')) return 'link'; + if (tag === 'button' || (tag === 'input' && ['button', 'submit', 'reset', 'image'].includes(type))) return 'button'; + if (tag === 'textarea' || element.isContentEditable || (tag === 'input' && !['button', 'submit', 'reset', 'image', 'checkbox', 'radio', 'hidden', 'file'].includes(type))) return 'textbox'; + if (tag === 'input' && type === 'checkbox') return 'checkbox'; + if (tag === 'input' && type === 'radio') return 'radio'; + if (tag === 'select') return element.multiple || element.size > 1 ? 'listbox' : 'combobox'; + if (tag === 'option') return 'option'; + if (/^h[1-6]$/.test(tag)) return 'heading'; + if (tag === 'img') return 'img'; + if (tag === 'ul' || tag === 'ol') return 'list'; + if (tag === 'li') return 'listitem'; + if (tag === 'table') return 'table'; + if (tag === 'tr') return 'row'; + if (tag === 'td') return 'cell'; + if (tag === 'th') return 'columnheader'; + return ''; + }; + const accessibleName = element => { + const ariaLabel = element.getAttribute('aria-label'); + if (ariaLabel) return normalize(ariaLabel); + const labelledBy = element.getAttribute('aria-labelledby'); + if (labelledBy) { + const text = labelledBy.split(/\s+/).map(id => document.getElementById(id)?.textContent || '').join(' '); + if (normalize(text)) return normalize(text); + } + if (element.labels?.length) return normalize(Array.from(element.labels).map(label => label.textContent).join(' ')); + if (element.tagName === 'IMG' && element.alt) return normalize(element.alt); + if (element.tagName === 'INPUT' && ['button', 'submit', 'reset'].includes((element.type || '').toLowerCase())) return normalize(element.value); + return normalize(element.getAttribute('title') || element.textContent || ''); + }; + + if (kind === 'testid') { + const attribute = testIdAttribute || 'data-testid'; + return elements.find(element => element.getAttribute(attribute) === value) || null; + } + if (kind === 'label') { + for (const element of elements) { + if (element.tagName !== 'LABEL' || !matches(element.textContent, value)) continue; + if (element.control) return element.control; + const nested = element.querySelector('input,textarea,select,[contenteditable="true"]'); + if (nested) return nested; + } + return elements.find(element => element.labels?.length && Array.from(element.labels).some(label => matches(label.textContent, value))) || null; + } + if (kind === 'role') { + return elements.find(element => implicitRole(element) === String(role || '').toLowerCase() && (!name || matches(accessibleName(element), name))) || null; + } + if (kind === 'text') { + const candidates = elements.filter(element => matches(element.innerText || element.textContent, value)); + return candidates.find(element => !Array.from(element.children).some(child => matches(child.innerText || child.textContent, value))) || candidates[0] || null; + } + return null; +}` + +// ParseSelector converts CSS/XPath and AIScan semantic locator syntax into the +// argument map used by nuclei headless actions. +// +// Supported semantic syntax: +// - text=Sign in +// - label=Email +// - testid=submit +// - role=button[name="Sign in"] +func ParseSelector(raw string) map[string]string { + raw = strings.TrimSpace(raw) + if xpath, ok := strings.CutPrefix(raw, "xpath:"); ok { + return map[string]string{"by": "xpath", "xpath": xpath} + } + for _, prefix := range []struct { + prefix string + by string + key string + }{ + {"text=", "text", "text"}, + {"label=", "label", "label"}, + {"testid=", "testid", "testid"}, + } { + if value, ok := strings.CutPrefix(raw, prefix.prefix); ok { + return map[string]string{"by": prefix.by, prefix.key: unquoteSelectorValue(value)} + } + } + if rest, ok := strings.CutPrefix(raw, "role="); ok { + args := map[string]string{"by": "role"} + role, attrs, _ := strings.Cut(rest, "[") + args["role"] = strings.TrimSpace(role) + attrs = strings.TrimSuffix(attrs, "]") + for _, attr := range strings.Split(attrs, "][") { + key, value, found := strings.Cut(attr, "=") + if found { + args[strings.TrimSpace(key)] = unquoteSelectorValue(value) + } + } + return args + } + return map[string]string{"selector": raw} +} + +func unquoteSelectorValue(value string) string { + value = strings.TrimSpace(value) + if len(value) >= 2 { + first, last := value[0], value[len(value)-1] + if (first == '\'' && last == '\'') || (first == '"' && last == '"') { + return value[1 : len(value)-1] + } + } + return value +} + +// FindElement resolves a CLI selector with the same semantics used by replay. +func FindElement(page *rod.Page, selector string, timeout time.Duration) (*rod.Element, error) { + if strings.TrimSpace(selector) == "" { + return nil, fmt.Errorf("empty selector") + } + return ElementBy(page, ParseSelector(selector), timeout) +} + +// ElementBy resolves a nuclei action selector. In addition to nuclei's +// CSS/XPath/regex/search forms, AIScan supports role, label, text, and testid. +func ElementBy(page *rod.Page, data map[string]string, timeout time.Duration) (*rod.Element, error) { + if timeout <= 0 { + timeout = defaultActionTimeout + } + page = page.Timeout(timeout) + by := strings.ToLower(strings.TrimSpace(data["by"])) + switch by { + case "x", "xpath": + xpath := data["xpath"] + if xpath == "" { + return nil, fmt.Errorf("xpath selector required") + } + return page.ElementX(xpath) + case "js": + if data["js"] == "" { + return nil, fmt.Errorf("js selector required") + } + return page.ElementByJS(rod.Eval(data["js"])) + case "r", "regex": + return page.ElementR(data["selector"], data["regex"]) + case "search": + result, err := page.Search(data["query"]) + if err != nil { + return nil, err + } + if result.First == nil { + return nil, fmt.Errorf("no element found for query: %s", data["query"]) + } + return result.First, nil + case "role", "label", "text", "testid": + value := data[by] + if by == "role" { + value = data["name"] + } + return page.ElementByJS(rod.Eval( + semanticSelectorJS, + by, + data["role"], + data["name"], + value, + strings.EqualFold(data["exact"], "true"), + data["testid-attribute"], + )) + default: + selector := data["selector"] + if selector == "" { + return nil, fmt.Errorf("no selector provided") + } + return page.Element(selector) + } +} diff --git a/pkg/imageutil/optimize.go b/pkg/imageutil/optimize.go new file mode 100644 index 00000000..87c53a08 --- /dev/null +++ b/pkg/imageutil/optimize.go @@ -0,0 +1,143 @@ +package imageutil + +import ( + "bytes" + "fmt" + "image" + "image/jpeg" + "image/png" + "io" + + "golang.org/x/image/draw" + _ "golang.org/x/image/webp" +) + +const ( + MaxDimension = 2000 + // Keeps inline media below a 4 MiB ProtoJSON WebSocket frame after base64 + // expansion and envelope overhead. Larger media travels by URI/file chunks. + MaxPayloadBytes = 2_500_000 +) + +var JPEGQualities = []int{85, 70, 55, 40} + +type Optimized struct { + MimeType string + Data []byte + OrigW int + OrigH int + FinalW int + FinalH int +} + +func Optimize(r io.Reader, srcMime string) (*Optimized, error) { + raw, err := io.ReadAll(r) + if err != nil { + return nil, err + } + if srcMime == "image/gif" { + return passthrough(raw, srcMime) + } + img, _, err := image.Decode(bytes.NewReader(raw)) + if err != nil { + return passthrough(raw, srcMime) + } + bounds := img.Bounds() + origW, origH := bounds.Dx(), bounds.Dy() + img = ResizeIfNeeded(img, origW, origH) + final := img.Bounds() + data, mime, err := pickSmallestEncoding(img) + if err != nil { + return nil, err + } + return &Optimized{ + MimeType: mime, + Data: data, + OrigW: origW, + OrigH: origH, + FinalW: final.Dx(), + FinalH: final.Dy(), + }, nil +} + +func OptimizeImage(img image.Image) (*Optimized, error) { + var raw bytes.Buffer + if err := png.Encode(&raw, img); err != nil { + return nil, err + } + return Optimize(bytes.NewReader(raw.Bytes()), "image/png") +} + +func passthrough(raw []byte, mime string) (*Optimized, error) { + if len(raw) > MaxPayloadBytes { + return nil, fmt.Errorf("image too large after encoding (%d bytes, max %d)", len(raw), MaxPayloadBytes) + } + return &Optimized{MimeType: mime, Data: raw}, nil +} + +func ResizeIfNeeded(img image.Image, w, h int) image.Image { + if w <= MaxDimension && h <= MaxDimension { + return img + } + var newW, newH int + if w > h { + newW = MaxDimension + newH = h * MaxDimension / w + } else { + newH = MaxDimension + newW = w * MaxDimension / h + } + if newW < 1 { + newW = 1 + } + if newH < 1 { + newH = 1 + } + dst := image.NewRGBA(image.Rect(0, 0, newW, newH)) + draw.CatmullRom.Scale(dst, dst.Bounds(), img, img.Bounds(), draw.Over, nil) + return dst +} + +func pickSmallestEncoding(img image.Image) ([]byte, string, error) { + pngData := EncodePNG(img) + jpegData := EncodeJPEG(img, JPEGQualities[0]) + best, mime := pngData, "image/png" + if len(jpegData) < len(best) { + best, mime = jpegData, "image/jpeg" + } + if len(best) <= MaxPayloadBytes { + return best, mime, nil + } + for _, quality := range JPEGQualities[1:] { + jpegData = EncodeJPEG(img, quality) + if len(jpegData) <= MaxPayloadBytes { + return jpegData, "image/jpeg", nil + } + } + bounds := img.Bounds() + w, h := bounds.Dx(), bounds.Dy() + for w > 1 && h > 1 { + w = max(1, w*3/4) + h = max(1, h*3/4) + dst := image.NewRGBA(image.Rect(0, 0, w, h)) + draw.CatmullRom.Scale(dst, dst.Bounds(), img, img.Bounds(), draw.Over, nil) + jpegData = EncodeJPEG(dst, JPEGQualities[0]) + if len(jpegData) <= MaxPayloadBytes { + return jpegData, "image/jpeg", nil + } + } + return nil, "", fmt.Errorf("cannot compress image to fit %d byte limit", MaxPayloadBytes) +} + +func EncodePNG(img image.Image) []byte { + var buf bytes.Buffer + enc := &png.Encoder{CompressionLevel: png.BestCompression} + _ = enc.Encode(&buf, img) + return buf.Bytes() +} + +func EncodeJPEG(img image.Image, quality int) []byte { + var buf bytes.Buffer + _ = jpeg.Encode(&buf, img, &jpeg.Options{Quality: quality}) + return buf.Bytes() +} diff --git a/pkg/imageutil/optimize_test.go b/pkg/imageutil/optimize_test.go new file mode 100644 index 00000000..b6e3ba7a --- /dev/null +++ b/pkg/imageutil/optimize_test.go @@ -0,0 +1,42 @@ +package imageutil + +import ( + "bytes" + "image" + "testing" +) + +func TestOptimizeImageResizesToPayloadBounds(t *testing.T) { + img := image.NewRGBA(image.Rect(0, 0, 4000, 1000)) + optimized, err := OptimizeImage(img) + if err != nil { + t.Fatal(err) + } + if optimized.OrigW != 4000 || optimized.OrigH != 1000 { + t.Fatalf("original dimensions = %dx%d", optimized.OrigW, optimized.OrigH) + } + if optimized.FinalW != MaxDimension || optimized.FinalH != 500 { + t.Fatalf("final dimensions = %dx%d, want %dx500", optimized.FinalW, optimized.FinalH, MaxDimension) + } + if len(optimized.Data) == 0 || len(optimized.Data) > MaxPayloadBytes { + t.Fatalf("optimized payload size = %d", len(optimized.Data)) + } +} + +func TestOptimizePassesThroughUnknownImageData(t *testing.T) { + raw := []byte("not-an-image") + optimized, err := Optimize(bytes.NewReader(raw), "application/octet-stream") + if err != nil { + t.Fatal(err) + } + if optimized.MimeType != "application/octet-stream" || !bytes.Equal(optimized.Data, raw) { + t.Fatalf("unexpected passthrough result: %+v", optimized) + } +} + +func TestOptimizeRejectsOversizedPassthrough(t *testing.T) { + _, err := Optimize(bytes.NewReader(make([]byte, MaxPayloadBytes+1)), "image/gif") + if err == nil { + t.Fatal("expected oversized passthrough error") + } +} diff --git a/pkg/node/agent.go b/pkg/node/agent.go index d12f564b..7da956f5 100644 --- a/pkg/node/agent.go +++ b/pkg/node/agent.go @@ -22,6 +22,9 @@ func RunWebSocket(ctx context.Context, option *cfg.Option, logger telemetry.Logg } func runRemoteAgent(ctx context.Context, option *cfg.Option, logger telemetry.Logger) error { + if err := resolveRemoteAgentURLs(option); err != nil { + return err + } nodeID, err := webNodeID(option) if err != nil { return err @@ -118,6 +121,16 @@ func runRemoteAgent(ctx context.Context, option *cfg.Option, logger telemetry.Lo return err } +func resolveRemoteAgentURLs(option *cfg.Option) error { + if option == nil { + return fmt.Errorf("web node configuration is required") + } + if err := cfg.ResolveAgentServerURLs(option); err != nil { + return fmt.Errorf("resolve remote agent URLs: %w", err) + } + return nil +} + // --------------------------------------------------------------------------- // chatAgentHandler implements the connection's upload and config-reload hooks. // AOP core/command messages are dispatched by rt.HandleEnvelope directly. diff --git a/pkg/node/agent_test.go b/pkg/node/agent_test.go index b69bd159..76f7e945 100644 --- a/pkg/node/agent_test.go +++ b/pkg/node/agent_test.go @@ -22,3 +22,31 @@ func TestWebNodeID(t *testing.T) { t.Fatal("expected missing node_id error") } } + +func TestResolveRemoteAgentURLsDerivesEmbeddedIOA(t *testing.T) { + option := &cfg.Option{ + AgentOptions: cfg.AgentOptions{ServerURL: "http://token@127.0.0.1:18080"}, + } + if err := resolveRemoteAgentURLs(option); err != nil { + t.Fatal(err) + } + if option.ServerURL != "http://token@127.0.0.1:18080" { + t.Fatalf("server URL = %q", option.ServerURL) + } + if option.IOAURL != "http://token@127.0.0.1:18080/ioa" { + t.Fatalf("IOA URL = %q, want same-origin embedded endpoint", option.IOAURL) + } +} + +func TestResolveRemoteAgentURLsPreservesIndependentIOA(t *testing.T) { + option := &cfg.Option{ + AgentOptions: cfg.AgentOptions{ServerURL: "http://token@127.0.0.1:18080"}, + IOAOptions: cfg.IOAOptions{IOAURL: "http://ioa-token@127.0.0.1:18765"}, + } + if err := resolveRemoteAgentURLs(option); err != nil { + t.Fatal(err) + } + if option.IOAURL != "http://ioa-token@127.0.0.1:18765" { + t.Fatalf("independent IOA URL = %q", option.IOAURL) + } +} diff --git a/pkg/node/proto_connection.go b/pkg/node/proto_connection.go index eb5293d4..0a417ba4 100644 --- a/pkg/node/proto_connection.go +++ b/pkg/node/proto_connection.go @@ -5,11 +5,15 @@ import ( "context" "errors" "fmt" + "io" + "mime" "net/http" + "net/url" "os" "os/exec" "path/filepath" "runtime" + "runtime/debug" "strconv" "strings" "sync" @@ -102,6 +106,7 @@ func connectGenerated(ctx context.Context, cc connectionConfig) error { if logger == nil { logger = telemetry.NopLogger() } + cc.Logger = logger attempt := 0 for { if ctx.Err() != nil { @@ -452,6 +457,12 @@ func handleAgentToolMessage(ctx context.Context, cc connectionConfig, envelope * taskCtx, taskCancel := context.WithCancel(ctx) trackOperation(operationsMu, operations, operationID, taskCancel) go func() { + defer func() { + if recovered := recover(); recovered != nil { + cc.Logger.Errorf("tool operation panic operation_id=%s tool=%s panic=%v\n%s", operationID, request.Call.Name, recovered, debug.Stack()) + fail("tool operation failed unexpectedly") + } + }() defer finishOperation(operationsMu, operations, operationID, taskCancel) event, err := runner.ExecuteToolRequest(taskCtx, operationID, request, cc.Registry, cc.Progress) if err != nil { @@ -597,10 +608,113 @@ func fileRead(req *filepb.ReadRequest, base string) fileResultValue { if req == nil || req.Path == "" { return fileResultValue{result: result, err: fmt.Errorf("file path is required")} } - data, err := os.ReadFile(resolveFileRPCPath(base, req.Path)) - result.Data = data - result.Size = int64(len(data)) - return fileResultValue{result: result, err: err} + requestPath, offset, limit, compat, err := normalizeFileReadRequest(req) + if err != nil { + return fileResultValue{result: result, err: err} + } + result.Path = requestPath + if offset < 0 { + return fileResultValue{result: result, err: fmt.Errorf("file offset cannot be negative")} + } + if limit < 0 { + return fileResultValue{result: result, err: fmt.Errorf("file read limit cannot be negative")} + } + path := resolveFileRPCPath(base, requestPath) + file, err := os.Open(path) + if err != nil { + return fileResultValue{result: result, err: err} + } + defer file.Close() + info, err := file.Stat() + if err != nil { + return fileResultValue{result: result, err: err} + } + if info.IsDir() { + return fileResultValue{result: result, err: fmt.Errorf("file path is a directory")} + } + if offset > info.Size() { + return fileResultValue{result: result, err: fmt.Errorf("file offset %d exceeds size %d", offset, info.Size())} + } + result.Filename = info.Name() + result.Size = info.Size() + result.Offset = offset + if limit == 0 { + data, readErr := io.ReadAll(file) + result.Data = data + result.Offset = 0 + result.Eof = readErr == nil + result.MediaType = detectFileMediaType(path, data) + finalizeCompatFileResult(result, compat) + return fileResultValue{result: result, err: readErr} + } + readLimit := min(int64(limit), int64(maxFileReadChunkBytes)) + remaining := info.Size() - offset + if readLimit > remaining { + readLimit = remaining + } + data := make([]byte, int(readLimit)) + n, readErr := file.ReadAt(data, offset) + if readErr != nil && readErr != io.EOF { + return fileResultValue{result: result, err: readErr} + } + result.Data = data[:n] + result.Eof = offset+int64(n) >= info.Size() + result.MediaType = detectFileMediaType(path, result.Data) + finalizeCompatFileResult(result, compat) + return fileResultValue{result: result} +} + +const maxFileReadChunkBytes int32 = 1 << 20 + +func detectFileMediaType(path string, data []byte) string { + if value := mime.TypeByExtension(strings.ToLower(filepath.Ext(path))); value != "" { + return value + } + if len(data) > 0 { + return http.DetectContentType(data) + } + return "application/octet-stream" +} + +func normalizeFileReadRequest(req *filepb.ReadRequest) (path string, offset int64, limit int32, compat bool, err error) { + parsed, parseErr := url.Parse(req.GetPath()) + if parseErr != nil || parsed.Scheme != "aop-range" || parsed.Host != "read" { + return req.GetPath(), req.GetOffset(), req.GetLimit(), false, nil + } + query := parsed.Query() + offset, err = strconv.ParseInt(query.Get("offset"), 10, 64) + if err != nil { + return "", 0, 0, true, fmt.Errorf("invalid range offset") + } + if rawLimit := query.Get("limit"); rawLimit != "" { + parsedLimit, limitErr := strconv.ParseInt(rawLimit, 10, 32) + if limitErr != nil { + return "", 0, 0, true, fmt.Errorf("invalid range limit") + } + limit = int32(parsedLimit) + } + path = query.Get("path") + if path == "" { + return "", 0, 0, true, fmt.Errorf("range file path is required") + } + return path, offset, limit, true, nil +} + +func finalizeCompatFileResult(result *filepb.Result, compat bool) { + if !compat || result == nil { + return + } + value := &url.URL{Scheme: "aop-range", Host: "result"} + query := value.Query() + query.Set("path", result.Path) + query.Set("offset", strconv.FormatInt(result.Offset, 10)) + if result.Eof { + query.Set("eof", "1") + } + value.RawQuery = query.Encode() + result.Path = value.String() + result.Offset = 0 + result.Eof = false } func fileWrite(req *filepb.WriteRequest, base string) fileResultValue { diff --git a/pkg/node/proto_connection_test.go b/pkg/node/proto_connection_test.go index aa237a99..cc7e09ed 100644 --- a/pkg/node/proto_connection_test.go +++ b/pkg/node/proto_connection_test.go @@ -1,17 +1,69 @@ package node import ( + "bytes" "context" + "strings" + "sync" + "testing" + "time" + + aop "github.com/chainreactors/aiscan/aop" execpb "github.com/chainreactors/aiscan/aop/exec" filepb "github.com/chainreactors/aiscan/aop/file" + toolpb "github.com/chainreactors/aiscan/aop/tool" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" protobuf "google.golang.org/protobuf/proto" "os" "path/filepath" "runtime" - "testing" ) +func TestToolOperationPanicIsReportedAndCleanedUp(t *testing.T) { + var logs bytes.Buffer + logger := telemetry.NewLogger(telemetry.LogConfig{Debug: true, Output: &logs}) + operations := make(map[string]context.CancelFunc) + var operationsMu sync.Mutex + failure := make(chan *aop.ProtocolError, 1) + send := func(_ string, message protobuf.Message) { + protocol := message.(*aop.ProtocolMessage) + if protocol.GetEvent() != nil { + panic("send event boom") + } + if value := protocol.GetProtocolError(); value != nil { + failure <- value + } + } + arguments, _ := aop.JSONValue(map[string]any{}) + request := &toolpb.Call{Call: &aop.ToolCall{Id: "op-panic", Name: "missing", Arguments: arguments}} + handleAgentToolMessage( + context.Background(), + connectionConfig{Registry: commands.NewRegistry(), Logger: logger}, + &aop.Envelope{Id: "op-panic"}, + &toolpb.ProtocolMessage{Message: &toolpb.ProtocolMessage_Call{Call: request}}, + send, &operationsMu, operations, + ) + + select { + case got := <-failure: + if got.Code != "OPERATION_FAILED" || !strings.Contains(got.Message, "unexpectedly") { + t.Fatalf("failure = %+v", got) + } + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for operation failure") + } + operationsMu.Lock() + _, tracked := operations["op-panic"] + operationsMu.Unlock() + if tracked { + t.Fatal("panicking operation was not cleaned up") + } + if got := logs.String(); !strings.Contains(got, "send event boom") || !strings.Contains(got, "op-panic") { + t.Fatalf("panic log = %s", got) + } +} + func TestExecRequestCompletesWithOutput(t *testing.T) { command := "printf hello" if runtime.GOOS == "windows" { @@ -98,6 +150,44 @@ func TestNativeFileRPCsResolveRelativeToRuntimeWorkdir(t *testing.T) { } } +func TestFileReadReturnsBoundedChunks(t *testing.T) { + base := t.TempDir() + path := filepath.Join(base, "capture.mp4") + data := bytes.Repeat([]byte("frame"), 300_000) + if err := os.WriteFile(path, data, 0o644); err != nil { + t.Fatal(err) + } + first := fileRead(&filepb.ReadRequest{Path: path, Limit: 256 * 1024}, base) + if first.err != nil { + t.Fatal(first.err) + } + if first.result.Offset != 0 || first.result.Eof || len(first.result.Data) != 256*1024 || first.result.Size != int64(len(data)) { + t.Fatalf("first chunk = %+v, bytes=%d", first.result, len(first.result.Data)) + } + if first.result.MediaType != "video/mp4" { + t.Fatalf("media type = %q, want video/mp4", first.result.MediaType) + } + joined := append([]byte(nil), first.result.Data...) + offset := int64(len(joined)) + for { + next := fileRead(&filepb.ReadRequest{Path: path, Offset: offset, Limit: maxFileReadChunkBytes + 1}, base) + if next.err != nil { + t.Fatal(next.err) + } + if next.result.Offset != offset || len(next.result.Data) > int(maxFileReadChunkBytes) { + t.Fatalf("chunk offset=%d bytes=%d, want offset=%d max=%d", next.result.Offset, len(next.result.Data), offset, maxFileReadChunkBytes) + } + joined = append(joined, next.result.Data...) + offset += int64(len(next.result.Data)) + if next.result.Eof { + break + } + } + if !bytes.Equal(joined, data) { + t.Fatalf("joined bytes = %d, want %d", len(joined), len(data)) + } +} + func TestUploadWritesAbsolutePath(t *testing.T) { const filename = "aiscan_test_upload_probe.txt" const body = "codex public proof\nkey=appImage/probe" diff --git a/pkg/runner/app.go b/pkg/runner/app.go index b4c9cee1..cc2292b5 100644 --- a/pkg/runner/app.go +++ b/pkg/runner/app.go @@ -78,6 +78,10 @@ func NewApp(ctx context.Context, rc ApplicationConfig) (*App, error) { logger = a.Logger() a.Hooks = hooks.New() a.Hooks.SetErrorSink(func(he *hooks.HandlerError) { + if len(he.Stack) > 0 { + a.Logger().Errorf("hook panic kind=%s source=%s panic=%v\n%s", he.Kind, he.Source, he.Panic, he.Stack) + return + } a.Logger().Warnf("hook failed kind=%s source=%s error=%q", he.Kind, he.Source, he.Err) }) @@ -375,13 +379,27 @@ func initCoreCommands(rc ApplicationConfig, llmProvider agent.Provider, skillSto Events: events, } plan := capability.Select(capability.Options{ - Groups: []string{"core", "arsenal", "search", "browser"}, + Groups: linkedToolGroups(), OptionalTools: rc.Tools.OptionalTools, }) commands.BuildPlan(plan, deps, cmdReg) + cmdReg.SetLogger(logger) return cmdReg } +func linkedToolGroups() []string { + seen := make(map[string]bool) + var groups []string + for _, descriptor := range capability.All() { + if descriptor.Kind != capability.KindTool || descriptor.Group == "" || seen[descriptor.Group] { + continue + } + seen[descriptor.Group] = true + groups = append(groups, descriptor.Group) + } + return groups +} + func executeRegistryCommand(ctx context.Context, reg *commands.CommandRegistry, commandLine string, timeout time.Duration) (string, error) { tool, ok := reg.GetTool("bash") if !ok { @@ -466,7 +484,9 @@ func (a *App) InitIOA(ctx context.Context, ioa IOAConfig) error { if ioa.AutoRegister { if err := client.EnsureRegistered(ctx, ioa.NodeName, "", ioa.NodeMeta); err != nil { a.Logger().Warnf("ioa registration pending: %s", err) - go a.retryIOARegistration(ctx, client, ioa) + telemetry.SafeGo("ioa-registration-retry", func() { + a.retryIOARegistration(ctx, client, ioa) + }) return nil } } diff --git a/pkg/runner/application_config.go b/pkg/runner/application_config.go index ff324d80..afb05737 100644 --- a/pkg/runner/application_config.go +++ b/pkg/runner/application_config.go @@ -45,7 +45,7 @@ type ToolConfig struct { BashTimeout int TavilyKeys string PlaywrightSession string - OptionalTools []string // optional tool groups to enable (e.g. "search", "browser") + OptionalTools []string // optional tool groups to enable } type IOAConfig struct { diff --git a/pkg/runner/ioa_safety.go b/pkg/runner/ioa_safety.go new file mode 100644 index 00000000..49229258 --- /dev/null +++ b/pkg/runner/ioa_safety.go @@ -0,0 +1,20 @@ +package runner + +import "reflect" + +// isNilIOADependency handles Go's typed-nil interface case. IOA clients are +// commonly stored behind protocol interfaces; assigning a nil *Client to one +// of those interfaces makes the interface itself non-nil and a plain +// dependency == nil check is therefore insufficient. +func isNilIOADependency(dependency any) bool { + if dependency == nil { + return true + } + value := reflect.ValueOf(dependency) + switch value.Kind() { + case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice: + return value.IsNil() + default: + return false + } +} diff --git a/pkg/runner/ioa_safety_test.go b/pkg/runner/ioa_safety_test.go new file mode 100644 index 00000000..465dc892 --- /dev/null +++ b/pkg/runner/ioa_safety_test.go @@ -0,0 +1,28 @@ +package runner + +import ( + "context" + "testing" + "time" + + inboxpkg "github.com/chainreactors/aiscan/agent/inbox" + ioaclient "github.com/chainreactors/ioa/client" +) + +func TestSubscribeIOASpaceTypedNilStreamReturns(t *testing.T) { + var concrete *ioaclient.Client + var stream ioaclient.StreamAPI = concrete + done := make(chan struct{}) + go func() { + subscribeIOASpace(context.Background(), stream, "space-1", "node-1", func(inboxpkg.Message) error { + return nil + }, nil) + close(done) + }() + + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("typed-nil IOA stream did not return") + } +} diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index 72dfc0e4..4c65dded 100644 --- a/pkg/runner/runner.go +++ b/pkg/runner/runner.go @@ -340,7 +340,7 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L Run: loop.Run, }, "loop") - if rt.app.IOAStreamClient != nil && option.Space != "" { + if !isNilIOADependency(rt.app.IOAStreamClient) && option.Space != "" { nodeID := "" if rt.app.IOAClient != nil { nodeID = rt.app.IOAClient.NodeID() @@ -351,7 +351,9 @@ func NewAgentRuntime(ctx context.Context, option *cfg.Option, logger telemetry.L } else { ioaCtx, cancel := context.WithCancel(ctx) ioaCancel = cancel - go subscribeIOASpace(ioaCtx, rt.app.IOAStreamClient, spaceInfo.ID, nodeID, rt.pushAsync, logger) + telemetry.SafeGo("ioa-space-subscription", func() { + subscribeIOASpace(ioaCtx, rt.app.IOAStreamClient, spaceInfo.ID, nodeID, rt.pushAsync, logger) + }) } } @@ -740,6 +742,12 @@ func scannerCommandSupportsDebug(name string) bool { // --------------------------------------------------------------------------- func subscribeIOASpace(ctx context.Context, stream ioaclient.StreamAPI, spaceID, nodeID string, push func(inboxpkg.Message) error, logger telemetry.Logger) { + if ctx == nil || isNilIOADependency(stream) || spaceID == "" || push == nil { + return + } + if logger == nil { + logger = telemetry.NopLogger() + } for attempt := 0; ctx.Err() == nil; attempt++ { msgs, errs, cancel, err := stream.Subscribe(ctx, spaceID) if err != nil { diff --git a/pkg/runner/runner_test.go b/pkg/runner/runner_test.go index af5e6450..86739035 100644 --- a/pkg/runner/runner_test.go +++ b/pkg/runner/runner_test.go @@ -161,6 +161,9 @@ func TestClearRotatesToAnEmptyContinuationSession(t *testing.T) { if _, err := run.Wait(); err != nil { t.Fatal(err) } + if runtime.sessionRunActive(session.ID()) { + t.Fatal("Run.Wait returned before the active run registration was released") + } oldID := session.ID() var events []*aop.Event diff --git a/pkg/runner/runtime_session.go b/pkg/runner/runtime_session.go index 246cfdf4..c55f89b0 100644 --- a/pkg/runner/runtime_session.go +++ b/pkg/runner/runtime_session.go @@ -1130,7 +1130,6 @@ func (s *sessionState) startRun(ctx context.Context, input RunInput) (*Run, erro emitter := &turnEmitter{sessionID: s.id, turnID: turnID, agentName: s.agentName, emitter: s.runtime.sessionEvents} op := &sessionOperation{ execute: func(runCtx context.Context) { - defer s.runtime.releaseRun(run) s.inbox.setActive(true) emitter.start() result, runErr := s.executeRun(runCtx, turnID, input) @@ -1149,17 +1148,16 @@ func (s *sessionState) startRun(ctx context.Context, input RunInput) (*Run, erro } emitter.end(runResult, runErr) s.inbox.setActive(false) - run.finish(runResult, runErr) + s.runtime.finishRun(run, runResult, runErr) }, reject: func(err error) { - defer s.runtime.releaseRun(run) result := RunResult{Stop: agent.StopReasonCanceled} if !errors.Is(err, context.Canceled) { result.Stop = agent.StopReasonError } emitter.start() emitter.end(result, err) - run.finish(result, err) + s.runtime.finishRun(run, result, err) }, } if err := s.admit(runCtx, op); err != nil { @@ -1349,13 +1347,26 @@ func (rt *AgentRuntime) releaseRun(run *Run) { if run == nil { return } + rt.unregisterRun(run) + rt.operations.Done() +} + +func (rt *AgentRuntime) finishRun(run *Run, result RunResult, err error) { + if run == nil { + return + } + rt.unregisterRun(run) + run.finish(result, err) + rt.operations.Done() +} + +func (rt *AgentRuntime) unregisterRun(run *Run) { run.cancel() rt.mu.Lock() if rt.runs[run.turnID] == run { delete(rt.runs, run.turnID) } rt.mu.Unlock() - rt.operations.Done() } func (rt *AgentRuntime) providerSnapshot() (agent.Provider, string, telemetry.Logger) { diff --git a/pkg/runner/subagent_handoff.go b/pkg/runner/subagent_handoff.go index 29d95be8..ccecb84b 100644 --- a/pkg/runner/subagent_handoff.go +++ b/pkg/runner/subagent_handoff.go @@ -17,7 +17,7 @@ import ( ) func subscribeIOAHandoffContext(ctx context.Context, bus *eventbus.Bus[*aop.Event], client protocols.ClientAPI, spaceName string, logger telemetry.Logger) func() { - if bus == nil || client == nil || spaceName == "" { + if bus == nil || isNilIOADependency(client) || spaceName == "" { return func() {} } if ctx == nil { @@ -43,7 +43,7 @@ func subscribeIOAHandoffContext(ctx context.Context, bus *eventbus.Bus[*aop.Even r.logger.Warnf("ioa handoff queue full, dropping %s", aop.Kind(event)) } }) - go r.run(ctx) + telemetry.SafeGo("ioa-handoff", func() { r.run(ctx) }) return func() { cancel() unsub() @@ -80,6 +80,9 @@ func (r *ioaHandoffRecorder) run(ctx context.Context) { case <-ctx.Done(): return case event := <-r.events: + if event == nil { + continue + } switch event.Payload.(type) { case *aop.Event_SessionStarted: r.onSessionStart(event) @@ -182,6 +185,9 @@ func (r *ioaHandoffRecorder) onTurnEnd(event *aop.Event) { } func (r *ioaHandoffRecorder) send(phase, status string, state *handoffState, title, message, refID string) (string, error) { + if r == nil || isNilIOADependency(r.client) { + return "", fmt.Errorf("IOA client is not configured") + } ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() spaceID, err := r.resolveSpace(ctx) diff --git a/pkg/runner/subagent_handoff_test.go b/pkg/runner/subagent_handoff_test.go index cce59855..beb5dd83 100644 --- a/pkg/runner/subagent_handoff_test.go +++ b/pkg/runner/subagent_handoff_test.go @@ -9,6 +9,7 @@ import ( aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/eventbus" types "github.com/chainreactors/aiscan/pkg/types" + ioaclient "github.com/chainreactors/ioa/client" "github.com/chainreactors/ioa/protocols" ) @@ -179,3 +180,23 @@ func TestIOAHandoffIgnoresNonDelegationSessions(t *testing.T) { time.Sleep(10 * time.Millisecond) } } + +func TestIOAHandoffTypedNilClientIsDisabled(t *testing.T) { + var concrete *ioaclient.Client + var client protocols.ClientAPI = concrete + if !isNilIOADependency(client) { + t.Fatal("typed-nil IOA client was treated as configured") + } + + bus := eventbus.New[*aop.Event]() + cancel := subscribeIOAHandoffContext(context.Background(), bus, client, "test", nil) + defer cancel() + + start := handoffEvent(t, "child-session", "worker", &aop.Event{Payload: &aop.Event_SessionStarted{SessionStarted: &aop.SessionStarted{ + ParentSessionId: "parent-session", ParentToolCallId: "spawn-typed-nil", + }}}) + if err := types.SetDelegation(start, &types.DelegationDetail{Task: "inspect target", AgentName: "worker"}); err != nil { + t.Fatal(err) + } + bus.Emit(start) +} diff --git a/pkg/runner/tool_call.go b/pkg/runner/tool_call.go index 0d721552..fdb535bb 100644 --- a/pkg/runner/tool_call.go +++ b/pkg/runner/tool_call.go @@ -21,19 +21,6 @@ type ToolExecutor interface { ExecuteTool(context.Context, string, string) (*tool.Result, error) } -// toolResolver is an optional executor capability exposing the concrete tool -// catalog, so the transport can assert execution capabilities on the resolved -// tool instead of its name. *commands.CommandRegistry implements it. -type toolResolver interface { - GetTool(name string) (tool.Tool, bool) -} - -// foregroundTool is implemented by tools that run a command in the foreground -// with streaming output, bypassing the agent-facing auto-background behavior. -type foregroundTool interface { - RunForegroundTool(context.Context, string, commands.BashExecOptions) (*tool.Result, error) -} - // ExecuteToolRequest runs one canonical AOP tool call against the executor // and wraps the outcome as a ToolResult event correlated to operationID. func ExecuteToolRequest(ctx context.Context, operationID string, request *toolpb.Call, executor ToolExecutor, progressBus *eventbus.Bus[*toolpb.Progress]) (*aop.Event, error) { @@ -79,22 +66,18 @@ func executeCall(ctx context.Context, executor ToolExecutor, call *aop.ToolCall, if len(arguments) == 0 { arguments = []byte("{}") } - if resolver, ok := executor.(toolResolver); ok { - if resolved, ok := resolver.GetTool(call.Name); ok { - if fg, ok := resolved.(foregroundTool); ok { - args, err := tool.ParseArgs[commands.BashArgs](string(arguments)) - if err != nil { - return nil, err - } - progress := newProgressStreamer(progressBus, call.Name, callID) - result, err := fg.RunForegroundTool(ctx, args.Command, commands.BashExecOptions{ - Timeout: time.Duration(args.Timeout) * time.Second, - OnOutput: progress.Write, - }) - progress.Flush() - return result, err - } + if registry, ok := executor.(*commands.CommandRegistry); ok && call.Name == "bash" { + args, err := tool.ParseArgs[commands.BashArgs](string(arguments)) + if err != nil { + return nil, err } + progress := newProgressStreamer(progressBus, call.Name, callID) + result, err := registry.ExecuteBashForeground(ctx, args.Command, commands.BashExecOptions{ + Timeout: time.Duration(args.Timeout) * time.Second, + OnOutput: progress.Write, + }) + progress.Flush() + return result, err } return executor.ExecuteTool(ctx, call.Name, string(arguments)) } diff --git a/pkg/runner/tool_call_test.go b/pkg/runner/tool_call_test.go index baf95e00..a735923b 100644 --- a/pkg/runner/tool_call_test.go +++ b/pkg/runner/tool_call_test.go @@ -1,6 +1,7 @@ package runner import ( + "bytes" "context" "errors" "strings" @@ -10,6 +11,7 @@ import ( aop "github.com/chainreactors/aiscan/aop" toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/eventbus" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/core/tool" "github.com/chainreactors/aiscan/pkg/commands" ) @@ -134,3 +136,32 @@ func TestExecuteToolRequestForeground(t *testing.T) { t.Fatalf("result = %+v", result) } } + +type panicForegroundBash struct{ recordingBash } + +func (*panicForegroundBash) RunForegroundTool(context.Context, string, commands.BashExecOptions) (*tool.Result, error) { + panic("foreground boom") +} + +func TestExecuteToolRequestForegroundPanicIsReturnedWithoutStack(t *testing.T) { + registry := commands.NewRegistry() + var logs bytes.Buffer + registry.SetLogger(telemetry.NewLogger(telemetry.LogConfig{Debug: true, Output: &logs})) + registry.RegisterTool(&panicForegroundBash{}) + + event, err := ExecuteToolRequest(context.Background(), "task-panic", toolRequest(t, "task-panic", "bash", map[string]any{"command": "echo test"}), registry, nil) + if err != nil { + t.Fatal(err) + } + result := event.GetToolResult() + text := tool.ResultText(result) + if !result.IsError || !strings.Contains(text, "task-panic") { + t.Fatalf("result = %+v", result) + } + if strings.Contains(text, "foreground boom") || strings.Contains(text, "goroutine") { + t.Fatalf("tool result leaked panic details: %q", text) + } + if got := logs.String(); !strings.Contains(got, "foreground boom") || !strings.Contains(got, "goroutine") { + t.Fatalf("panic log = %s", got) + } +} diff --git a/pkg/web/service/scan.go b/pkg/web/service/scan.go index a30da281..a9373223 100644 --- a/pkg/web/service/scan.go +++ b/pkg/web/service/scan.go @@ -10,11 +10,13 @@ import ( "net" "net/url" "os" + "runtime/debug" "strings" aop "github.com/chainreactors/aiscan/aop" "github.com/chainreactors/aiscan/core/output" coretool "github.com/chainreactors/aiscan/core/tool" + "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/pkg/commands" types "github.com/chainreactors/aiscan/pkg/types" managementapi "github.com/chainreactors/aiscan/pkg/web/api" @@ -167,8 +169,9 @@ func (s *Service) runScan(runCtx context.Context, scanID string) { }() defer func() { if recovered := recover(); recovered != nil { + telemetry.GlobalLogs().Errorf("scan panic scan_id=%s panic=%v\n%s", scanID, recovered, debug.Stack()) if scan, err := s.store.Get(context.Background(), scanID); err == nil { - _, _ = s.failScan(scan, fmt.Sprintf("scan runtime panic: %v", recovered)) + _, _ = s.failScan(scan, "scan failed unexpectedly") } } }() diff --git a/skills/aiscan/SKILL.md b/skills/aiscan/SKILL.md index f543df87..8b494096 100644 --- a/skills/aiscan/SKILL.md +++ b/skills/aiscan/SKILL.md @@ -25,6 +25,7 @@ Core agent tools: - `bash`: run shell commands and pseudo-commands (see below). - `web_search`: search the web for CVEs, advisories, exploits, and documentation. - `fetch`: fetch and read a specific URL. +- `record` (Windows/Linux full builds): capture desktop or visible application-window screenshots and H.264 recordings. It accepts HWND/X11 Window IDs or resolves a PID to its main visible window. ## User Tool Restrictions @@ -69,7 +70,7 @@ Available only when they appear in the runtime pseudo-command list: - `cyberhub`: search fingerprints and POC templates. Key: `cyberhub search --finger `. Reference: `aiscan://skills/aiscan/okf/runtime/search.md`. - `tmux`: session management. Key: `tmux ls`, `tmux capture-pane -t `, `tmux kill-session -t `. Reference: `aiscan://skills/aiscan/okf/runtime/tmux.md`. - `proxy`: proxy nodes and proxied execution. Key: `proxy `, `proxy auto `. Reference: `aiscan://skills/aiscan/okf/runtime/proxy.md`. -- `ioa`: multi-agent collaboration via shared message spaces — `ioa space `, `ioa send`, `ioa read --all`, and `ioa send checkpoint`. Reference: `aiscan://skills/ioa/SKILL.md`; wire-protocol formats live in the ioa module skills (`ioa://skills//SKILL.md`). +- `ioa`: multi-agent collaboration via shared message spaces — `ioa space `, `ioa send`, `ioa read --all`, and `ioa send checkpoint`. Reference: `aiscan://skills/ioa/SKILL.md`; wire-protocol formats live in the ioa module skills (`ioa://skills//SKILL.md`). Publish vulnerability discoveries per `aiscan://skills/aiscan/okf/runtime/ioa-finding.md`. ## Fingerprint → POC Workflow diff --git a/skills/aiscan/okf/easm/katana.md b/skills/aiscan/okf/easm/katana.md index 1f2857f4..a09c0330 100644 --- a/skills/aiscan/okf/easm/katana.md +++ b/skills/aiscan/okf/easm/katana.md @@ -35,8 +35,30 @@ katana -u https://target.com -d 2 -jsonl katana -u https://target.com -f qurl katana -u https://target.com -d 3 -jc -jsonl katana -list urls.txt -d 2 -jc -timeout 60 + +# Rendered browser crawling +katana -u https://target.com -hl -d 2 -jsonl +katana -u https://target.com -hh -d 2 -jsonl +katana -u https://target.com -hl --chrome-ws-url ws://127.0.0.1:9222/devtools/browser/ ``` +## Browser Modes and Reuse + +- Standard Katana crawling does not launch a browser. `-jc` parses JavaScript responses but does not render the application. +- `-hl` runs pure headless crawling and captures browser requests, dynamic navigation, forms, and rendered interactions. +- `-hh` combines HTTP crawling with browser rendering. Prefer `-hl` when browser network events and SPA navigation must be emitted as results. +- The full scan profile's `katana_deep` capability uses pure headless crawling; `katana_crawl` remains the lower-cost standard crawler. + +AIScan resolves a browser in this order: + +1. Katana `--chrome-ws-url` (`-cwu`) for an explicitly managed running process. +2. Katana `--system-chrome-path` (`-scp`) for an explicitly selected executable. +3. `AISCAN_BROWSER_PATH` shared by AIScan Playwright, nuclei headless replay, and Katana. +4. Installed Chrome, Chromium, or Edge in PATH or the standard OS install locations. +5. Rod's existing browser cache, with its first-use download only when the cache is absent. + +Automatic reuse means the executable is shared while each engine starts an isolated process/profile. AIScan does not automatically attach to a user's running browser. Use `--chrome-ws-url` only when process-level reuse is intentional. + ## Useful Filters - `-f qurl` — only output URLs that contain query parameters diff --git a/skills/aiscan/okf/easm/playwright.md b/skills/aiscan/okf/easm/playwright.md index 63847e9d..35dde136 100644 --- a/skills/aiscan/okf/easm/playwright.md +++ b/skills/aiscan/okf/easm/playwright.md @@ -269,7 +269,7 @@ playwright unroute # Remove all request inte ## Recording (nuclei headless template codegen) -Record browser interactions as a nuclei-compatible headless YAML template. This is aiscan's equivalent of Playwright's `codegen` — but outputs nuclei headless YAML instead of test scripts. +Record successful browser commands as a nuclei-shaped headless YAML template. This is aiscan's codegen workflow: it records CLI operations after they succeed and emits declarative browser actions instead of Node.js test code. ### Enable recording ```bash @@ -280,7 +280,7 @@ playwright open http://target.com/login --session s1 --record playwright record s1 --start ``` -When `--record` is active, every interaction command (click, fill, press, select-option, wait-for, eval, etc.) is automatically captured as a nuclei headless action. +When `--record` is active, supported interaction, navigation, extraction, storage, cookie, and wait commands are captured automatically. `fill` records a clear-then-input operation, while `type` appends. `press` keeps both the target selector and key expression (for example `Control+A` or `Shift+Enter`). ### Export recorded template ```bash @@ -301,7 +301,7 @@ playwright template poc.yaml http://other-target.com playwright template poc.yaml http://other-target.com --payload username=admin --payload password=test ``` -The generated YAML is standard nuclei headless format — it can also be used with neutron or nuclei directly. +The generated YAML uses the nuclei headless schema. Templates containing only the upstream core actions remain portable to compatible nuclei/neutron runners. Actions marked as AIScan extensions require `playwright template` in AIScan; upstream nuclei does not know those action names. ### Recording workflow example ```bash @@ -325,21 +325,99 @@ playwright template interaction.yaml http://target3.com/search |---|---| | `open --record` (initial) | `navigate` with `{{BaseURL}}` | | `click` | `click` | -| `fill` / `type` | `text` | -| `press` | `keyboard` | -| `select-option` | `select` | +| `fill` | `text` with `clear: "true"` | +| `type` | `text` (append) | +| `press` | `keyboard` with selector and `keys` | +| `select-option` | `select` with `selected: "true"` | +| `set-input-files` / `upload` | `files` | | `eval` | `script` | | `wait-for --stable` | `waitstable` | | `wait-for --idle` | `waitidle` | | `wait-for ` | `waitvisible` | +| `wait-for-url/request/response` | `waiturl` / `waitrequest` / `waitresponse` (AIScan) | | `text-content` / `inner-text` | `extract` (with auto-generated name) | +| `content` / `inner-html` | `extract` with `target: html` | | `get-attribute` | `extract` (target=attribute) | -| `screenshot` | `screenshot` | +| `input-value`, `url`, `title` | `extract` with the corresponding target | +| `is-visible/hidden/checked/enabled/disabled` | `assert` preserving the observed boolean state (AIScan) | +| `screenshot` | `screenshot`, including `--selector` | | `set-extra-headers` | `setheader` (one per header) | -| `dialog --arm` | `waitdialog` | -| `hover` / `dblclick` / `reload` | `script` (JS fallback) | +| `hover`, `dblclick`, `focus`, `blur` | same-named AIScan action | +| `check`, `uncheck` | idempotent same-named AIScan action | +| `dispatch-event` | `dispatch` (AIScan) | +| `set-viewport` | `setviewport` (AIScan) | +| `reload`, `go-back`, `go-forward` | `reload`, `goback`, `goforward` (AIScan) | +| `set-content` | `setcontent` (AIScan) | +| local/session storage set/delete/clear | `storage` (AIScan) | +| local/session storage get/list | `extract` with `target: storage` | +| cookie set/delete/clear | `cookie` (AIScan) | +| cookie get/list | `extract` with `target: cookie` | +| `dialog --arm`, `dialog-accept`, `dialog-dismiss` | non-blocking `dialog` handler | -URLs are automatically templatized: the session's base origin is replaced with `{{BaseURL}}`. XPath selectors (`xpath:...`) are preserved as `by: xpath`. +URLs are automatically templatized: the session's base origin is replaced with `{{BaseURL}}`. + +### Selector vocabulary + +Live CLI operations and recorded template replay use the same selector resolver. Semantic selectors traverse the document and open shadow roots. + +| syntax | meaning | +|---|---| +| `input[name=email]` | CSS selector | +| `xpath://button[@type='submit']` | XPath selector | +| `text=Sign in` | visible text substring | +| `label=Email` | form control associated with a label | +| `testid=submit` | exact `data-testid` value | +| `role=button[name="Sign in"]` | implicit/explicit ARIA role and accessible name | + +Recorded semantic selectors are stored as structured action args (`by`, `role`, `name`, `label`, `testid`, and so on), so replay does not fall back to `document.querySelector`. In hand-written YAML, add `exact: "true"` for exact semantic text/name matching, or `testid-attribute` to override `data-testid`. + +### AIScan headless extensions + +These actions are available to `playwright template` in addition to the upstream nuclei-compatible core set. + +| action | important args | behavior | +|---|---|---| +| `dblclick`, `hover`, `focus`, `blur` | selector args | Native Rod element interaction | +| `check`, `uncheck` | selector args | Set the desired checked state; repeated replay is safe | +| `dispatch` | selector, `event`, optional JSON `detail` | Dispatch `Event` or `CustomEvent` | +| `setviewport` | `width`, `height`, optional `device-scale-factor` | Change viewport metrics | +| `waiturl` | `url`, optional `match` | Wait for current URL | +| `waitrequest`, `waitresponse` | `url`, optional `method`, `match`, `timeout` | Match captured browser traffic | +| `storage` | `storage`, `operation`, `key`, `value` | Set/delete/clear localStorage or sessionStorage | +| `cookie` | `operation`, `name`, `value`, optional URL/domain/path flags | Set/delete/clear cookies | +| `assert` | `type`, selector/value-specific args, optional `match` | Verify visible DOM, value, attribute, URL, title, storage, or cookie state | +| `scroll` | `x`, `y`, `steps` | Mouse-wheel scrolling | +| `drag` | source selector args, `target` | Drag the source element to a target selector | +| `reload`, `goback`, `goforward` | optional `timeout` | Browser history navigation followed by stability wait | +| `setcontent` | `html` | Replace the current document content | + +String waits and assertions accept `match: contains`, `equals`, or `regex`. Boolean assertion types are `visible`, `hidden`, `checked`, `unchecked`, `enabled`, and `disabled`; value assertions include `text`, `value`, `attribute`, `url`, `title`, `storage`, and `cookie`. + +Example extension steps: + +```yaml +- action: text + args: + by: label + label: Email + value: user@example.com + clear: "true" +- action: check + args: + by: testid + testid: terms +- action: assert + args: + by: role + role: button + name: Continue + type: visible +- action: waitresponse + args: + url: /api/session + method: POST + match: contains +``` ## Headless Template Execution @@ -349,7 +427,7 @@ Run a nuclei-compatible headless YAML template against a target URL. Shares the playwright template [--payload key=value ...] ``` -Templates support the full nuclei headless action set (29 action types), DSL expressions (`{{rand_int()}}`, `{{replace()}}`, etc.), payload iteration (sniper/pitchfork/clusterbomb), template variables, matchers, and extractors. +Templates support the 29-action nuclei-compatible core plus the AIScan extensions above, DSL expressions (`{{rand_int()}}`, `{{replace()}}`, etc.), payload iteration (sniper/pitchfork/clusterbomb), template variables, matchers, and extractors. ```bash # Run a recorded template @@ -451,8 +529,10 @@ Use browser automation when evidence depends on rendered DOM, user interaction, - setTimeout/setInterval acceleration (0.1x factor, disable with `--no-speed-up`) - Console messages are auto-captured from session open — retrieve with `console `. - Sessions persist until explicitly closed — the agent is responsible for calling `playwright close`. -- Chromium is automatically downloaded on first launch if not found. -- Selectors may be CSS or `xpath:` — interaction commands accept both. +- Browser discovery uses `AISCAN_BROWSER_PATH` first, then installed Chrome, Chromium, or Edge in the system PATH and standard OS install locations. Rod's cached/downloaded Chromium is used only when neither is available. Katana uses the same discovery policy. +- System-browser reuse means reusing the executable while AIScan launches an isolated managed process/profile. To control an already running browser process, use `attach --cdp ` or `open --cdp `. +- `playwright template` injects the already connected Playwright Rod browser into the nuclei-compatible headless engine, so template replay does not launch or download a second browser. +- Selectors may be CSS, `xpath:`, `text=...`, `label=...`, `testid=...`, or `role=[name="..."]`; live commands and template replay share the resolver. ## Related concepts diff --git a/skills/aiscan/okf/easm/scan.md b/skills/aiscan/okf/easm/scan.md index 15e15bd7..4222e4e8 100644 --- a/skills/aiscan/okf/easm/scan.md +++ b/skills/aiscan/okf/easm/scan.md @@ -59,7 +59,8 @@ Notes: - `quick` uses gogo ports `all`, spray check/finger, spray crawl depth 2, weakpass checks, and fingerprint-based POC checks. - `full` uses gogo ports `-` and adds spray plugins (common/bak/active) plus spray default-dictionary probing; crawl depth remains 2. -- Full builds additionally add katana crawling: `katana_crawl` in quick/full and `katana_deep` in full. +- Full builds additionally add Katana crawling: `katana_crawl` in quick/full uses the standard HTTP engine, while `katana_deep` in full uses pure headless rendering and emits browser-only requests and SPA navigation. +- `katana_deep` shares AIScan's browser discovery order: `AISCAN_BROWSER_PATH`, installed Chrome/Chromium/Edge, then Rod's cache/download fallback. - Spray web capabilities run with recon enabled in both profiles. - `--verify=` triggers in-pipeline AI verification for loots at or above the specified priority threshold (low, medium, high, critical). Only loots meeting the threshold are sent to a verify sub-agent. - `--sniper` asks an LLM agent to perform fingerprint vulnerability intelligence. diff --git a/skills/aiscan/okf/runtime/index.md b/skills/aiscan/okf/runtime/index.md index 01237e9b..332cd2e8 100644 --- a/skills/aiscan/okf/runtime/index.md +++ b/skills/aiscan/okf/runtime/index.md @@ -11,6 +11,7 @@ This bundle organizes aiscan's runtime mechanism tool documentation as concept f - [ioa space](ioa-space.md) — IOA space selection and discovery - [ioa send](ioa-send.md) — IOA messages and checkpoints - [ioa read](ioa-read.md) — IOA inbox and thread reading +- [ioa finding](ioa-finding.md) — publishing vulnerability discoveries as IOA checkpoints - [search](search.md) — cyberhub fingerprint/POC search - [fetch](fetch.md) — URL content retrieval and focused extraction - [loop](loop.md) — recurring agent task scheduling diff --git a/skills/aiscan/okf/runtime/ioa-finding.md b/skills/aiscan/okf/runtime/ioa-finding.md new file mode 100644 index 00000000..fd87d220 --- /dev/null +++ b/skills/aiscan/okf/runtime/ioa-finding.md @@ -0,0 +1,73 @@ +--- +type: Tool Playbook +title: ioa finding +description: Publish vulnerability discoveries to the current IOA space as natural-language checkpoints, making them observable and reviewable by other nodes. +tags: [runtime, collaboration, ioa, finding] +status: stable +--- + +# ioa finding — Publishing Vulnerability Discoveries + +When a scan produces a vulnerability worth acting on, publish it to the current +IOA space as a `checkpoint` message. IOA persists the message, any node can +watch the stream with `ioa read --listen`, and the checkpoint status forms a +review trail. + +Requires an IOA-bound session (a space already selected via +[ioa space](ioa-space.md) or configured by the runner). If `ioa send` reports +"no space joined", skip publishing — findings still live in the local report. + +## When to publish + +Your judgment. Publish: + +- **Verified vulnerabilities** — active probing confirmed the issue. +- **Verified weak credentials** — confirmed login on a real service. +- **High-value leads needing human review** — you cannot complete verification + but the evidence is strong enough that a human should look. + +Do NOT publish: + +- Fingerprints, service banners, routine web artifacts. +- Sniper/CVE intelligence without exploitation evidence. +- Leads you already dismissed (strikethrough / not confirmed). + +## How to publish + +``` +ioa send checkpoint --kind finding \ + --title "" \ + --target "" \ + --status pending \ + --content "" +``` + +- `--status pending` — verified or high-severity, asks for human review. +- `--status info` — intelligence-level, no review requested. + +Write the content natural-language first: what it is, where, how you verified +it, and the impact. End with a reference line linking the durable records: + +``` +result_id: · evidence: mitm://flows/17, neutron PoC shiro-rce.yaml +``` + +The `result_id` is the finding's only identity — content-addressed from the +artifact data, so the same vulnerability re-scanned yields the same id and +anyone holding the artifact can recompute it. The full record lives at +`findings/.md` in the local report bundle (see the report +reference); the IOA message is the observable notification. + +## Review loop + +A human or peer node reviews the checkpoint and replies with another +`ioa send checkpoint --kind finding` on the same target carrying +`--status confirmed` or `--status dismissed`. Treat that status as the +disposition record; mirror it into the local `findings/.md` +frontmatter when you next touch the report. + +## Related concepts + +- Send mechanics and flags: [ioa send](ioa-send.md) +- Watching the stream: [ioa read](ioa-read.md) +- Space selection: [ioa space](ioa-space.md) diff --git a/skills/aiscan/reference/report.md b/skills/aiscan/reference/report.md index 965e443e..14295690 100644 --- a/skills/aiscan/reference/report.md +++ b/skills/aiscan/reference/report.md @@ -11,10 +11,12 @@ Write the report as a directory of markdown files, borrowing mechanisms from OKF / index.md # human-readable summary (the six sections below) findings/ - .md # one concept file per confirmed finding or noteworthy lead + .md # one concept file per confirmed finding or noteworthy lead ``` -Each `findings/.md` carries frontmatter: +A finding's only identity is its `result_id` — the content-addressed artifact id assigned at scan time, stable across re-scans of the same data. There is no separate finding id or slug; the filename is the id. + +Each `findings/.md` carries frontmatter: ```yaml --- @@ -77,7 +79,7 @@ Count confirmed vulnerabilities separately from unverified leads. Strikethrough List verified loots first. Unannotated scanner matches may appear only with "unverified scanner match" stated clearly. For each: -- **[target]** — vulnerability description, CVE if applicable, impact, verification status, link to findings/.md with the reproducible PoC +- **[target]** — vulnerability description, CVE if applicable, impact, verification status, link to findings/.md with the reproducible PoC ## Potential Risks (Unverified) @@ -110,3 +112,7 @@ Brief list so the reader knows what was checked and cleared. - Prioritize by severity: critical > high > medium. - Use plain markdown, no code fences around the report. - If no significant loots remain after applying verification filters, say so clearly. An honest "no confirmed vulnerabilities" is far more valuable than inflated severity. + +## Publishing to IOA + +When the session is IOA-bound, publish each confirmed finding to the current space as a `checkpoint` message (`--kind finding`, natural-language content) per `aiscan://skills/aiscan/okf/runtime/ioa-finding.md`. Cite the `result_id` in the message — it is the whole link: the IOA message is the observable notification, `findings/.md` is the complete record. Review replies on the checkpoint (confirmed/dismissed) are the disposition trail — mirror them into the finding frontmatter on the next report update. diff --git a/skills/embed_test.go b/skills/embed_test.go index 324a8e11..85b3ce30 100644 --- a/skills/embed_test.go +++ b/skills/embed_test.go @@ -124,6 +124,36 @@ func TestReadVirtual(t *testing.T) { } } +func TestIOAFindingConvention(t *testing.T) { + store, _ := LoadEmbeddedStore() + body, handled, err := store.ReadVirtualBody("aiscan://skills/aiscan/okf/runtime/ioa-finding.md") + if err != nil || !handled { + t.Fatalf("ReadVirtualBody(ioa-finding) handled=%v err=%v", handled, err) + } + if !strings.Contains(body, "--kind finding") || !strings.Contains(body, "result_id") { + t.Fatalf("ioa-finding convention missing checkpoint command or result_id:\n%s", body) + } + if strings.Contains(body, "finding-id") { + t.Fatal("ioa-finding must not reintroduce a separate finding-id") + } + + main, _, err := store.ReadVirtual("aiscan://skills/aiscan/SKILL.md") + if err != nil { + t.Fatalf("ReadVirtual(SKILL.md) error = %v", err) + } + if !strings.Contains(main, "okf/runtime/ioa-finding.md") { + t.Fatal("SKILL.md does not reference the ioa-finding convention") + } + + report, _, err := store.ReadVirtual("aiscan://skills/aiscan/reference/report.md") + if err != nil { + t.Fatalf("ReadVirtual(report.md) error = %v", err) + } + if !strings.Contains(report, "findings/.md") { + t.Fatal("report.md must name findings by result_id") + } +} + func TestLoadAllIncludesIOAModuleSkills(t *testing.T) { store, diags := LoadAll(nil) if len(diags) != 0 { diff --git a/test-skips.json b/test-skips.json index a185eedf..b8d59c9a 100644 --- a/test-skips.json +++ b/test-skips.json @@ -62,6 +62,13 @@ "category": "platform", "reason": "Validates Unix-only PTY and process lifecycle behavior." }, + { + "path": "cmd/aiscan/imports_full_test.go", + "format": "record is only linked on Windows and Linux", + "count": 1, + "category": "platform", + "reason": "The native recorder capability is linked only into Windows and Linux full builds." + }, { "path": "pkg/commands/bash_test.go", "format": "unix-only test", @@ -139,6 +146,13 @@ "category": "live_llm", "reason": "The IOA integration calls a live model endpoint and requires an explicit credential." }, + { + "path": "tools/katana/katana_test.go", + "format": "no system browser available; CI installs Chrome for this test", + "count": 1, + "category": "external_runtime", + "reason": "The Katana browser capability test requires a locally installed Chromium-family executable." + }, { "path": "tools/playwright/browser_test.go", "format": "no Chromium/Chrome found, skipping browser integration test", @@ -180,5 +194,12 @@ "count": 1, "category": "external_api", "reason": "The Hunter integration requires a user-supplied service credential." + }, + { + "path": "tools/scan/capability_katana_test.go", + "format": "no system browser available; CI installs Chrome for this test", + "count": 1, + "category": "external_runtime", + "reason": "The scan capability test requires a locally installed Chromium-family executable." } ] diff --git a/tools/imports_record_native.go b/tools/imports_record_native.go new file mode 100644 index 00000000..a4b435b8 --- /dev/null +++ b/tools/imports_record_native.go @@ -0,0 +1,5 @@ +//go:build full && record_ffmpeg && cgo && (windows || linux) + +package tools + +import _ "github.com/chainreactors/aiscan/tools/record" diff --git a/tools/katana/katana.go b/tools/katana/katana.go index daa850dc..bd2796ce 100644 --- a/tools/katana/katana.go +++ b/tools/katana/katana.go @@ -14,6 +14,7 @@ import ( aop "github.com/chainreactors/aiscan/aop" toolpb "github.com/chainreactors/aiscan/aop/tool" "github.com/chainreactors/aiscan/core/telemetry" + browserutil "github.com/chainreactors/aiscan/pkg/browser" "github.com/chainreactors/aiscan/pkg/commands" "github.com/chainreactors/aiscan/tools/toolargs" "github.com/projectdiscovery/goflags" @@ -139,6 +140,9 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any if options.Proxy == "" && c.Proxy != "" { options.Proxy = c.Proxy } + if err := configureBrowserOptions(options); err != nil { + return nil, fmt.Errorf("katana: %w", err) + } if err := validateOptions(options); err != nil { return nil, fmt.Errorf("katana: %w", err) @@ -216,6 +220,35 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any return nil, nil } +type browserDiscoverer func() (browserutil.Binary, error) + +func configureBrowserOptions(options *katanatypes.Options) error { + return configureBrowserOptionsWith(options, browserutil.Discover) +} + +func configureBrowserOptionsWith(options *katanatypes.Options, discover browserDiscoverer) error { + if options.ChromeWSUrl != "" { + return nil + } + if options.SystemChromePath != "" { + options.UseInstalledChrome = true + return nil + } + if !options.Headless && !options.HeadlessHybrid { + return nil + } + + binary, err := discover() + if err != nil { + return fmt.Errorf("browser discovery failed: %w", err) + } + if binary.Path != "" { + options.SystemChromePath = binary.Path + options.UseInstalledChrome = true + } + return nil +} + // readFlags replicates katana's cmd/katana/main.go readFlags() using goflags, // keeping CLI arguments 100% compatible with the upstream katana binary. func readFlags(args []string) (*katanatypes.Options, error) { diff --git a/tools/katana/katana_test.go b/tools/katana/katana_test.go new file mode 100644 index 00000000..61c5f190 --- /dev/null +++ b/tools/katana/katana_test.go @@ -0,0 +1,197 @@ +package katana + +import ( + "bytes" + "context" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + browserutil "github.com/chainreactors/aiscan/pkg/browser" + "github.com/chainreactors/aiscan/pkg/commands" + katanatypes "github.com/projectdiscovery/katana/pkg/types" +) + +func TestConfigureBrowserOptionsPriority(t *testing.T) { + tests := []struct { + name string + options katanatypes.Options + discovered browserutil.Binary + discoveryErr error + wantPath string + wantInstalled bool + wantCalls int + wantErr bool + }{ + { + name: "CDP endpoint takes precedence", + options: katanatypes.Options{Headless: true, ChromeWSUrl: "ws://127.0.0.1:9222/devtools/browser/test"}, + wantCalls: 0, + }, + { + name: "explicit Katana path takes precedence", + options: katanatypes.Options{Headless: true, SystemChromePath: "/explicit/chrome"}, + wantPath: "/explicit/chrome", + wantInstalled: true, + wantCalls: 0, + }, + { + name: "headless uses shared discovery", + options: katanatypes.Options{Headless: true}, + discovered: browserutil.Binary{Path: "/system/chromium", Source: browserutil.SourceSystem}, + wantPath: "/system/chromium", + wantInstalled: true, + wantCalls: 1, + }, + { + name: "hybrid uses shared discovery", + options: katanatypes.Options{HeadlessHybrid: true}, + discovered: browserutil.Binary{Path: "/system/edge", Source: browserutil.SourceSystem}, + wantPath: "/system/edge", + wantInstalled: true, + wantCalls: 1, + }, + { + name: "standard crawler does not need a browser", + options: katanatypes.Options{}, + wantCalls: 0, + }, + { + name: "missing browser preserves Rod fallback", + options: katanatypes.Options{Headless: true}, + wantCalls: 1, + }, + { + name: "discovery errors are returned", + options: katanatypes.Options{Headless: true}, + discoveryErr: errors.New("bad browser override"), + wantCalls: 1, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + calls := 0 + err := configureBrowserOptionsWith(&tt.options, func() (browserutil.Binary, error) { + calls++ + return tt.discovered, tt.discoveryErr + }) + if (err != nil) != tt.wantErr { + t.Fatalf("configureBrowserOptionsWith error = %v, wantErr %v", err, tt.wantErr) + } + if calls != tt.wantCalls { + t.Fatalf("discover calls = %d, want %d", calls, tt.wantCalls) + } + if tt.options.SystemChromePath != tt.wantPath { + t.Fatalf("SystemChromePath = %q, want %q", tt.options.SystemChromePath, tt.wantPath) + } + if tt.options.UseInstalledChrome != tt.wantInstalled { + t.Fatalf("UseInstalledChrome = %v, want %v", tt.options.UseInstalledChrome, tt.wantInstalled) + } + }) + } +} + +func TestE2EHeadlessReusesDiscoveredBrowser(t *testing.T) { + binary, err := browserutil.Discover() + if err != nil { + t.Fatalf("discover browser: %v", err) + } + if binary.Path == "" { + t.Skip("no system browser available; CI installs Chrome for this test") + } + t.Setenv(browserutil.PathEnv, binary.Path) + + const ( + sessionToken = "aiscan-session-42" + workspacePath = "/workspace/session-42?view=issues" + ) + var rootHits atomic.Int32 + var sessionHits atomic.Int32 + var authenticatedWorkspaceHits atomic.Int32 + + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/" { + http.NotFound(w, r) + return + } + rootHits.Add(1) + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprint(w, ` +AIScan Workspace Login +
Signing in...
+`) + }) + mux.HandleFunc("/api/session", func(w http.ResponseWriter, r *http.Request) { + sessionHits.Add(1) + if r.Method != http.MethodPost || r.Header.Get("X-CSRF-Token") != "browser-e2e" { + http.Error(w, "invalid session request", http.StatusForbidden) + return + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"token":%q,"next":%q}`, sessionToken, workspacePath) + }) + mux.HandleFunc("/workspace/session-42", func(w http.ResponseWriter, r *http.Request) { + cookie, err := r.Cookie("aiscan_session") + if err != nil || cookie.Value != sessionToken { + http.Error(w, "authentication required", http.StatusUnauthorized) + return + } + authenticatedWorkspaceHits.Add(1) + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprint(w, `

Security issues

Critical issue`) + }) + mux.HandleFunc("/issues/critical", func(w http.ResponseWriter, _ *http.Request) { + fmt.Fprint(w, "confirmed issue") + }) + srv := httptest.NewServer(mux) + defer srv.Close() + + options, err := readFlags([]string{"-u", srv.URL, "-hl", "-d", "2"}) + if err != nil { + t.Fatalf("parse headless options: %v", err) + } + if err := configureBrowserOptions(options); err != nil { + t.Fatalf("configure browser options: %v", err) + } + if options.SystemChromePath != binary.Path || !options.UseInstalledChrome { + t.Fatalf("Katana browser = (%q, installed=%v), want (%q, true)", options.SystemChromePath, options.UseInstalledChrome, binary.Path) + } + + ctx, cancel := context.WithTimeout(context.Background(), 75*time.Second) + defer cancel() + var output bytes.Buffer + _, err = New().Run(ctx, &commands.Execution{ + Args: []string{"-u", srv.URL, "-hl", "-d", "2", "-timeout", "15", "-ct", "45s", "-j"}, + Stdout: &output, + Stderr: &output, + }) + if err != nil { + t.Fatalf("Katana headless crawl failed: %v\noutput:\n%s", err, output.String()) + } + if authenticatedWorkspaceHits.Load() == 0 { + t.Fatalf("browser never reached the authenticated workspace (root=%d session=%d)\noutput:\n%s", rootHits.Load(), sessionHits.Load(), output.String()) + } + if !strings.Contains(output.String(), "/workspace/session-42") { + t.Fatalf("Katana output does not contain the browser-only workspace route\noutput:\n%s", output.String()) + } +} diff --git a/tools/neutron/neutron.go b/tools/neutron/neutron.go index ba6d1ddb..eda55e3d 100644 --- a/tools/neutron/neutron.go +++ b/tools/neutron/neutron.go @@ -61,21 +61,6 @@ type neutronFlags struct { Debug bool `long:"debug" description:"Enable debug logging"` } -type neutronResult struct { - Timestamp string `json:"timestamp,omitempty"` - Target string `json:"target"` - Matched bool `json:"matched"` - Template string `json:"template"` - Name string `json:"name,omitempty"` - Severity string `json:"severity,omitempty"` - Tags []string `json:"tags,omitempty"` - Fingers []string `json:"fingers,omitempty"` - Extracts []string `json:"extracts,omitempty"` - Request string `json:"request,omitempty"` - Response string `json:"response,omitempty"` - Error string `json:"error,omitempty"` -} - type neutronSummary struct { Targets int Templates int @@ -245,11 +230,11 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any if result.Error() != nil { summary.Errors++ } - record := neutronResultFromExecution(target, result) - results = append(results, result.TemplateResult(target)) + record := result.TemplateResult(target) + results = append(results, record) if record.Matched { summary.Matched++ - c.EmitArtifactCtx(ctx, "neutron", toolpb.ArtifactKindVuln, target, &record) + c.EmitArtifactCtx(ctx, "neutron", toolpb.ArtifactKindVuln, target, record) } if shouldPrintNeutronResult(record, flags) { line := formatNeutronResult(record, jsonOutput) @@ -424,7 +409,7 @@ func validateNeutronSeverities(groups ...[]string) error { return nil } -func shouldPrintNeutronResult(record neutronResult, flags neutronFlags) bool { +func shouldPrintNeutronResult(record *sdktypes.TemplateResult, flags neutronFlags) bool { if flags.AllResults { return true } @@ -434,34 +419,7 @@ func shouldPrintNeutronResult(record neutronResult, flags neutronFlags) bool { return record.Matched } -func neutronResultFromExecution(target string, result *sdkneutron.ExecuteResult) neutronResult { - record := neutronResult{ - Timestamp: time.Now().Format(time.RFC3339), - Target: target, - Matched: result.Matched(), - } - if tmpl := result.Template(); tmpl != nil { - record.Template = tmpl.Id - record.Name = tmpl.Info.Name - record.Severity = tmpl.Info.Severity - record.Tags = cleanTemplateTags(tmpl) - record.Fingers = append([]string(nil), tmpl.Fingers...) - } - if opResult := result.Value(); opResult != nil { - record.Extracts = append([]string(nil), opResult.OutputExtracts()...) - // The engine captures the exchange (protocols/http/request.go sets these - // on the operator result); dropping it here left every consumer unable to - // show what was actually sent, so a match had no reviewable evidence. - record.Request = opResult.Request - record.Response = opResult.Response - } - if err := result.Error(); err != nil { - record.Error = err.Error() - } - return record -} - -func formatNeutronResult(record neutronResult, jsonOutput bool) string { +func formatNeutronResult(record *sdktypes.TemplateResult, jsonOutput bool) string { if jsonOutput { data, err := json.Marshal(record) if err != nil { @@ -479,25 +437,17 @@ func formatNeutronResult(record neutronResult, jsonOutput bool) string { b.WriteString(status) b.WriteString("] ") b.WriteString(record.Target) - if record.Template != "" { + if record.TemplateID != "" { b.WriteString(" template=") - b.WriteString(record.Template) + b.WriteString(record.TemplateID) } if record.Severity != "" { b.WriteString(" severity=") b.WriteString(record.Severity) } - if record.Name != "" { + if record.TemplateName != "" { b.WriteString(" name=") - b.WriteString(strconv.Quote(record.Name)) - } - if len(record.Extracts) > 0 { - b.WriteString(" extracts=") - b.WriteString(strconv.Quote(strings.Join(record.Extracts, ","))) - } - if record.Error != "" { - b.WriteString(" error=") - b.WriteString(strconv.Quote(record.Error)) + b.WriteString(strconv.Quote(record.TemplateName)) } b.WriteByte('\n') return b.String() @@ -541,12 +491,12 @@ func renderTemplateList(selected []*templates.Template, jsonOutput bool) string if tmpl == nil { continue } - record := neutronResult{ - Template: tmpl.Id, - Name: tmpl.Info.Name, - Severity: tmpl.Info.Severity, - Tags: cleanTemplateTags(tmpl), - Fingers: append([]string(nil), tmpl.Fingers...), + record := map[string]any{ + "template_id": tmpl.Id, + "template_name": tmpl.Info.Name, + "severity": tmpl.Info.Severity, + "tags": cleanTemplateTags(tmpl), + "fingers": append([]string(nil), tmpl.Fingers...), } if jsonOutput { data, err := json.Marshal(record) diff --git a/tools/neutron/neutron_test.go b/tools/neutron/neutron_test.go index 4dedf5ef..43ab56a5 100644 --- a/tools/neutron/neutron_test.go +++ b/tools/neutron/neutron_test.go @@ -85,11 +85,11 @@ func TestCommandTemplateListSupportsNucleiStyleFlagsAndJSON(t *testing.T) { t.Fatalf("Execute() error = %v", err) } out := output.String() - var result neutronResult + var result map[string]any if err := json.Unmarshal([]byte(strings.TrimSpace(out)), &result); err != nil { t.Fatalf("json output = %q, error = %v", out, err) } - if result.Template != "critical-cve" || result.Severity != "critical" { + if result["template_id"] != "critical-cve" || result["severity"] != "critical" { t.Fatalf("result = %#v", result) } } @@ -175,9 +175,8 @@ func TestNeutronResultFromExecutionCarriesExchange(t *testing.T) { Response: "HTTP/1.1 200 OK\r\n\r\nidentifier=3.3M2.0", } op.Matched = true - record := neutronResultFromExecution( + record := (&sdkneutron.ExecuteResult{TypedResult: sdktypes.NewResult(true, nil, op)}).TemplateResult( "https://example.test", - &sdkneutron.ExecuteResult{TypedResult: sdktypes.NewResult(true, nil, op)}, ) if record.Request != op.Request { t.Fatalf("request not carried through: got %q, want %q", record.Request, op.Request) diff --git a/tools/playwright/browser.go b/tools/playwright/browser.go index 4365f4b7..1cc518f8 100644 --- a/tools/playwright/browser.go +++ b/tools/playwright/browser.go @@ -16,6 +16,7 @@ import ( "github.com/chainreactors/aiscan/core/telemetry" "github.com/chainreactors/aiscan/core/truncate" + browserutil "github.com/chainreactors/aiscan/pkg/browser" "github.com/chainreactors/aiscan/pkg/commands" "github.com/go-rod/rod" "github.com/go-rod/rod/lib/launcher" @@ -529,7 +530,7 @@ func (c *Command) Run(ctx context.Context, execution *commands.Execution) (_ any if err == nil && len(subArgs) > 0 { if sess, sessErr := c.getSession(subArgs[0]); sessErr == nil && sess.rec != nil { - recordCommand(sess, sub, subArgs) + recordCommandResult(sess, sub, subArgs, result) } } @@ -583,11 +584,12 @@ func (c *Command) getOrLaunchBrowser() (*rod.Browser, error) { Set("disable-dev-shm-usage"). Set("ignore-certificate-errors"). Set("allow-insecure-localhost") - // Prefer an installed browser when one is available. Rod otherwise - // enters its auto-download path, which is inappropriate for offline CI - // and can block even though launcher.LookPath already found Chromium. - if browserPath, ok := launcher.LookPath(); ok { - l = l.Bin(browserPath) + binary, err := browserutil.Discover() + if err != nil { + return nil, fmt.Errorf("playwright: browser discovery failed: %w", err) + } + if binary.Path != "" { + l = l.Bin(binary.Path) } c.proxyMu.RLock() diff --git a/tools/playwright/interact.go b/tools/playwright/interact.go index 342e9d9b..eb2c340e 100644 --- a/tools/playwright/interact.go +++ b/tools/playwright/interact.go @@ -12,6 +12,7 @@ import ( "strings" "time" + "github.com/chainreactors/aiscan/pkg/headless" "github.com/go-rod/rod" "github.com/go-rod/rod/lib/input" "github.com/go-rod/rod/lib/proto" @@ -1055,14 +1056,7 @@ func (c *Command) execType(ctx context.Context, args []string) (string, error) { // --------------------------------------------------------------------------- func findElement(page *rod.Page, selector string) (*rod.Element, error) { - selector = strings.TrimSpace(selector) - if selector == "" { - return nil, fmt.Errorf("empty selector") - } - if xpath, ok := strings.CutPrefix(selector, "xpath:"); ok { - return page.ElementX(xpath) - } - return page.Element(selector) + return headless.FindElement(page, selector, 0) } func selectOption(el *rod.Element, value string) error { diff --git a/tools/playwright/recorder.go b/tools/playwright/recorder.go index 6f6ce55b..354f9fee 100644 --- a/tools/playwright/recorder.go +++ b/tools/playwright/recorder.go @@ -8,6 +8,7 @@ import ( "fmt" "net/url" "os" + "strconv" "strings" "sync" @@ -105,6 +106,10 @@ func (r *recorder) generateTemplate(id, name string) *headless.Template { // recordCommand maps a playwright command invocation to a nuclei headless action // and appends it to the session's recorder. Returns true if the action was recorded. func recordCommand(sess *Session, cmd string, args []string) bool { + return recordCommandResult(sess, cmd, args, "") +} + +func recordCommandResult(sess *Session, cmd string, args []string, result string) bool { if sess.rec == nil { return false } @@ -132,10 +137,8 @@ func recordCommand(sess *Session, cmd string, args []string) bool { return false } ra = RecordedAction{ - Action: headless.ActionScript, - Args: map[string]string{ - "code": fmt.Sprintf(`document.querySelector(%q).dispatchEvent(new MouseEvent('dblclick', {bubbles: true}))`, sel), - }, + Action: headless.ActionDblClick, + Args: selectorArgs(sel), } case "fill": @@ -146,7 +149,7 @@ func recordCommand(sess *Session, cmd string, args []string) bool { value := strings.Join(args[2:], " ") ra = RecordedAction{ Action: headless.ActionTextInput, - Args: mergeMaps(selectorArgs(sel), map[string]string{"value": value}), + Args: mergeMaps(selectorArgs(sel), map[string]string{"value": value, "clear": "true"}), } case "type": @@ -167,7 +170,7 @@ func recordCommand(sess *Session, cmd string, args []string) bool { keys := strings.Join(args[2:], " ") ra = RecordedAction{ Action: headless.ActionKeyboard, - Args: map[string]string{"keys": keys}, + Args: mergeMaps(selectorArgs(args[1]), map[string]string{"keys": keys}), } case "select-option", "select": @@ -178,7 +181,7 @@ func recordCommand(sess *Session, cmd string, args []string) bool { value := strings.Join(args[2:], " ") ra = RecordedAction{ Action: headless.ActionSelectInput, - Args: mergeMaps(selectorArgs(sel), map[string]string{"value": value}), + Args: mergeMaps(selectorArgs(sel), map[string]string{"value": value, "selected": "true"}), } case "screenshot": @@ -192,10 +195,13 @@ func recordCommand(sess *Session, cmd string, args []string) bool { } else if args[i] == "--output" && i+1 < len(args) { i++ ra.Args["to"] = args[i] + } else if args[i] == "--selector" && i+1 < len(args) { + i++ + ra.Args = mergeMaps(ra.Args, selectorArgs(args[i])) } } - case "set-input-files": + case "set-input-files", "upload": if len(args) < 3 { return false } @@ -222,10 +228,8 @@ func recordCommand(sess *Session, cmd string, args []string) bool { return false } ra = RecordedAction{ - Action: headless.ActionScript, - Args: map[string]string{ - "code": fmt.Sprintf(`document.querySelector(%q).dispatchEvent(new MouseEvent('mouseover', {bubbles: true}))`, sel), - }, + Action: headless.ActionHover, + Args: selectorArgs(sel), } case "wait-for", "wait": @@ -247,7 +251,7 @@ func recordCommand(sess *Session, cmd string, args []string) bool { default: ra = RecordedAction{ Action: headless.ActionWaitVisible, - Args: map[string]string{"selector": target}, + Args: selectorArgs(target), } } @@ -269,20 +273,20 @@ func recordCommand(sess *Session, cmd string, args []string) bool { case "reload": ra = RecordedAction{ - Action: headless.ActionScript, - Args: map[string]string{"code": "window.location.reload()"}, + Action: headless.ActionReload, + Args: map[string]string{}, } case "go-back", "back": ra = RecordedAction{ - Action: headless.ActionScript, - Args: map[string]string{"code": "window.history.back()"}, + Action: headless.ActionGoBack, + Args: map[string]string{}, } case "go-forward", "forward": ra = RecordedAction{ - Action: headless.ActionScript, - Args: map[string]string{"code": "window.history.forward()"}, + Action: headless.ActionGoForward, + Args: map[string]string{}, } case "check": @@ -291,7 +295,7 @@ func recordCommand(sess *Session, cmd string, args []string) bool { return false } ra = RecordedAction{ - Action: headless.ActionClick, + Action: headless.ActionCheck, Args: selectorArgs(sel), } @@ -301,7 +305,7 @@ func recordCommand(sess *Session, cmd string, args []string) bool { return false } ra = RecordedAction{ - Action: headless.ActionClick, + Action: headless.ActionUncheck, Args: selectorArgs(sel), } @@ -322,16 +326,14 @@ func recordCommand(sess *Session, cmd string, args []string) bool { sel := args[1] eventType := args[2] ra = RecordedAction{ - Action: headless.ActionScript, - Args: map[string]string{ - "code": fmt.Sprintf(`document.querySelector(%q).dispatchEvent(new Event(%q, {bubbles: true}))`, sel, eventType), - }, + Action: headless.ActionDispatchEvent, + Args: mergeMaps(selectorArgs(sel), map[string]string{"event": eventType}), } case "dialog": if len(args) >= 2 && args[1] == "--arm" { ra = RecordedAction{ - Action: headless.ActionWaitDialog, + Action: headless.ActionDialog, Args: map[string]string{}, } } else { @@ -349,6 +351,17 @@ func recordCommand(sess *Session, cmd string, args []string) bool { Name: sanitizeName(sel), } + case "content", "inner-html", "html": + sel := "html" + if len(args) >= 2 { + sel = strings.Join(args[1:], " ") + } + ra = RecordedAction{ + Action: headless.ActionExtract, + Args: mergeMaps(selectorArgs(sel), map[string]string{"target": "html"}), + Name: sanitizeName(sel + "_html"), + } + case "get-attribute": if len(args) < 3 { return false @@ -372,14 +385,159 @@ func recordCommand(sess *Session, cmd string, args []string) bool { ra = RecordedAction{ Action: headless.ActionExtract, Args: mergeMaps(selectorArgs(sel), map[string]string{ - "target": "attribute", - "attribute": "value", + "target": "value", }), Name: sanitizeName(sel + "_value"), } case "set-viewport": - return false + if len(args) < 3 { + return false + } + ra = RecordedAction{ + Action: headless.ActionSetViewport, + Args: map[string]string{"width": args[1], "height": args[2]}, + } + + case "focus", "blur": + sel := extractSelector(args, 1) + if sel == "" { + return false + } + action := headless.ActionFocus + if cmd == "blur" { + action = headless.ActionBlur + } + ra = RecordedAction{Action: action, Args: selectorArgs(sel)} + + case "wait-for-url", "wait-for-request", "wait-for-response": + if len(args) < 2 { + return false + } + action := headless.ActionWaitURL + if cmd == "wait-for-request" { + action = headless.ActionWaitRequest + } else if cmd == "wait-for-response" { + action = headless.ActionWaitResponse + } + ra = RecordedAction{Action: action, Args: map[string]string{"url": strings.Join(args[1:], " ")}} + + case "set-content": + if len(args) < 2 { + return false + } + ra = RecordedAction{Action: headless.ActionSetContent, Args: map[string]string{"html": strings.Join(args[1:], " ")}} + + case "url", "title": + ra = RecordedAction{ + Action: headless.ActionExtract, + Args: map[string]string{"target": cmd}, + Name: cmd, + } + + case "is-visible", "is-hidden", "is-checked", "is-disabled", "is-enabled": + sel := extractSelector(args, 1) + if sel == "" { + return false + } + assertionType := strings.TrimPrefix(cmd, "is-") + if strings.HasSuffix(strings.TrimSpace(result), "= false") { + assertionType = map[string]string{ + "visible": "hidden", "hidden": "visible", + "checked": "unchecked", "disabled": "enabled", "enabled": "disabled", + }[assertionType] + } + ra = RecordedAction{ + Action: headless.ActionAssert, + Args: mergeMaps(selectorArgs(sel), map[string]string{"type": assertionType}), + } + + case "localstorage-set", "sessionstorage-set": + if len(args) < 3 { + return false + } + storageType := strings.TrimSuffix(cmd, "-set") + ra = RecordedAction{Action: headless.ActionStorage, Args: map[string]string{ + "storage": storageType, "operation": "set", "key": args[1], "value": strings.Join(args[2:], " "), + }} + + case "localstorage-delete", "sessionstorage-delete": + if len(args) < 2 { + return false + } + storageType := strings.TrimSuffix(cmd, "-delete") + ra = RecordedAction{Action: headless.ActionStorage, Args: map[string]string{ + "storage": storageType, "operation": "delete", "key": args[1], + }} + + case "localstorage-clear", "sessionstorage-clear": + storageType := strings.TrimSuffix(cmd, "-clear") + ra = RecordedAction{Action: headless.ActionStorage, Args: map[string]string{ + "storage": storageType, "operation": "clear", + }} + + case "localstorage-get", "sessionstorage-get": + if len(args) < 2 { + return false + } + storageType := strings.TrimSuffix(cmd, "-get") + ra = RecordedAction{Action: headless.ActionExtract, Args: map[string]string{ + "target": "storage", "storage": storageType, "key": args[1], + }, Name: sanitizeName(storageType + "_" + args[1])} + + case "localstorage-list", "sessionstorage-list": + storageType := strings.TrimSuffix(cmd, "-list") + ra = RecordedAction{Action: headless.ActionExtract, Args: map[string]string{ + "target": "storage", "storage": storageType, + }, Name: sanitizeName(storageType)} + + case "cookie-set": + if len(args) < 2 { + return false + } + recorded := false + for _, pair := range args[1:] { + name, value, ok := strings.Cut(pair, "=") + if !ok || name == "" { + continue + } + sess.rec.record(RecordedAction{Action: headless.ActionCookie, Args: map[string]string{ + "operation": "set", "name": name, "value": value, + }}) + recorded = true + } + return recorded + + case "cookie-delete": + if len(args) < 2 { + return false + } + ra = RecordedAction{Action: headless.ActionCookie, Args: map[string]string{ + "operation": "delete", "name": args[1], + }} + + case "cookie-clear": + ra = RecordedAction{Action: headless.ActionCookie, Args: map[string]string{"operation": "clear"}} + + case "cookie-get": + if len(args) < 2 { + return false + } + ra = RecordedAction{Action: headless.ActionExtract, Args: map[string]string{ + "target": "cookie", "name": args[1], + }, Name: sanitizeName("cookie_" + args[1])} + + case "cookie-list": + ra = RecordedAction{Action: headless.ActionExtract, Args: map[string]string{ + "target": "cookie", + }, Name: "cookies"} + + case "dialog-accept", "dialog-dismiss": + argsMap := map[string]string{"accept": strconv.FormatBool(cmd == "dialog-accept")} + if cmd == "dialog-accept" && len(args) >= 2 { + argsMap["prompt"] = strings.Join(args[1:], " ") + } + ra = RecordedAction{Action: headless.ActionDialog, Args: argsMap} default: return false @@ -503,11 +661,7 @@ func recordSave(sess *Session, path, id, name string) (string, error) { // selectorArgs converts a CSS/XPath selector string to nuclei action args. func selectorArgs(sel string) map[string]string { - sel = strings.TrimSpace(sel) - if xpath, ok := strings.CutPrefix(sel, "xpath:"); ok { - return map[string]string{"by": "xpath", "xpath": xpath} - } - return map[string]string{"selector": sel} + return headless.ParseSelector(sel) } // extractSelector extracts a selector from args starting at the given offset. diff --git a/tools/playwright/recorder_test.go b/tools/playwright/recorder_test.go index e75f5ee1..7542e522 100644 --- a/tools/playwright/recorder_test.go +++ b/tools/playwright/recorder_test.go @@ -5,6 +5,7 @@ package playwright import ( "bytes" "context" + "encoding/json" "fmt" "io" "net/http" @@ -12,6 +13,7 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" "github.com/chainreactors/aiscan/pkg/commands" @@ -149,14 +151,31 @@ func TestRecordCommandMapping(t *testing.T) { {"wait-for", []string{"test", "--stable"}, headless.ActionWaitStable}, {"wait-for", []string{"test", "--idle"}, headless.ActionWaitIdle}, {"wait-for", []string{"test", "#element"}, headless.ActionWaitVisible}, - {"hover", []string{"test", "#menu"}, headless.ActionScript}, - {"dblclick", []string{"test", "#item"}, headless.ActionScript}, - {"reload", []string{"test"}, headless.ActionScript}, - {"go-back", []string{"test"}, headless.ActionScript}, - {"dialog", []string{"test", "--arm"}, headless.ActionWaitDialog}, + {"hover", []string{"test", "#menu"}, headless.ActionHover}, + {"dblclick", []string{"test", "#item"}, headless.ActionDblClick}, + {"reload", []string{"test"}, headless.ActionReload}, + {"go-back", []string{"test"}, headless.ActionGoBack}, + {"go-forward", []string{"test"}, headless.ActionGoForward}, + {"dialog", []string{"test", "--arm"}, headless.ActionDialog}, {"text-content", []string{"test", "#result"}, headless.ActionExtract}, {"get-attribute", []string{"test", "a", "href"}, headless.ActionExtract}, {"inner-text", []string{"test", "#text"}, headless.ActionExtract}, + {"inner-html", []string{"test", "#markup"}, headless.ActionExtract}, + {"check", []string{"test", "#terms"}, headless.ActionCheck}, + {"uncheck", []string{"test", "#terms"}, headless.ActionUncheck}, + {"focus", []string{"test", "#email"}, headless.ActionFocus}, + {"blur", []string{"test", "#email"}, headless.ActionBlur}, + {"dispatch-event", []string{"test", "#form", "submit"}, headless.ActionDispatchEvent}, + {"set-viewport", []string{"test", "1280", "720"}, headless.ActionSetViewport}, + {"wait-for-url", []string{"test", "/done"}, headless.ActionWaitURL}, + {"wait-for-request", []string{"test", "/api/start"}, headless.ActionWaitRequest}, + {"wait-for-response", []string{"test", "/api/result"}, headless.ActionWaitResponse}, + {"set-content", []string{"test", "
ready
"}, headless.ActionSetContent}, + {"localstorage-set", []string{"test", "token", "abc"}, headless.ActionStorage}, + {"sessionstorage-delete", []string{"test", "draft"}, headless.ActionStorage}, + {"cookie-set", []string{"test", "sid=123"}, headless.ActionCookie}, + {"cookie-delete", []string{"test", "sid"}, headless.ActionCookie}, + {"is-visible", []string{"test", "#ready"}, headless.ActionAssert}, } for _, tt := range tests { @@ -175,6 +194,41 @@ func TestRecordCommandMapping(t *testing.T) { } } +func TestRecordCommandReplaySemantics(t *testing.T) { + sess := &Session{Name: "test", rec: newRecorder("https://example.com")} + + recordCommand(sess, "fill", []string{"test", "label=Email", "alice@example.com"}) + recordCommand(sess, "press", []string{"test", `role=button[name="Sign in"]`, "Shift+Enter"}) + recordCommand(sess, "wait-for", []string{"test", "testid=ready"}) + + actions := sess.rec.snapshot() + if len(actions) != 3 { + t.Fatalf("expected 3 actions, got %d", len(actions)) + } + if actions[0].Args["clear"] != "true" || actions[0].Args["by"] != "label" || actions[0].Args["label"] != "Email" { + t.Fatalf("fill did not preserve clear/label semantics: %#v", actions[0].Args) + } + if actions[1].Args["by"] != "role" || actions[1].Args["role"] != "button" || actions[1].Args["name"] != "Sign in" { + t.Fatalf("press did not preserve role selector: %#v", actions[1].Args) + } + if actions[1].Args["keys"] != "Shift+Enter" { + t.Fatalf("press keys = %q", actions[1].Args["keys"]) + } + if actions[2].Args["by"] != "testid" || actions[2].Args["testid"] != "ready" { + t.Fatalf("wait selector was not parsed semantically: %#v", actions[2].Args) + } +} + +func TestRecordCommandResultPreservesBooleanState(t *testing.T) { + sess := &Session{Name: "test", rec: newRecorder("https://example.com")} + if !recordCommandResult(sess, "is-visible", []string{"test", "#optional"}, "#optional visible = false") { + t.Fatal("is-visible result was not recorded") + } + if got := sess.rec.snapshot()[0].Args["type"]; got != "hidden" { + t.Fatalf("false is-visible result recorded as %q, want hidden", got) + } +} + func TestRecordCommandXPath(t *testing.T) { sess := &Session{ Name: "test", @@ -751,3 +805,256 @@ func TestIntegration_RecordRoundTrip(t *testing.T) { t.Errorf("expected template ID in output, got:\n%s", out) } } + +func TestIntegration_RecordExtendedRoundTrip(t *testing.T) { + skipIfNoBrowserR(t) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprint(w, ` + + + + + +`) + })) + defer srv.Close() + + workDir := t.TempDir() + cmd := New(workDir) + defer cmd.Close() + ctx := context.Background() + + recExecString(t, cmd, ctx, []string{"open", srv.URL, "--session", "extended", "--record", "--timeout", "10"}) + recExecString(t, cmd, ctx, []string{"fill", "extended", "label=Email", "alice@example.com"}) + recExecString(t, cmd, ctx, []string{"press", "extended", "label=Email", "End"}) + recExecString(t, cmd, ctx, []string{"check", "extended", "testid=terms"}) + recExecString(t, cmd, ctx, []string{"select-option", "extended", `role=combobox[name="Plan"]`, "pro"}) + recExecString(t, cmd, ctx, []string{"hover", "extended", `role=button[name="Continue"]`}) + recExecString(t, cmd, ctx, []string{"dblclick", "extended", `role=button[name="Continue"]`}) + recExecString(t, cmd, ctx, []string{"dispatch-event", "extended", "#continue", "aiscan"}) + recExecString(t, cmd, ctx, []string{"localstorage-set", "extended", "token", "abc123"}) + recExecString(t, cmd, ctx, []string{"cookie-set", "extended", "session=cookie-value"}) + recExecString(t, cmd, ctx, []string{"set-viewport", "extended", "1024", "768"}) + recExecString(t, cmd, ctx, []string{"is-checked", "extended", "testid=terms"}) + + templatePath := filepath.Join(workDir, "extended-roundtrip.yaml") + recExecString(t, cmd, ctx, []string{"record", "extended", "--save", templatePath, "--id", "extended-roundtrip"}) + data, err := os.ReadFile(templatePath) + if err != nil { + t.Fatal(err) + } + tmpl, err := headless.ParseTemplate(data) + if err != nil { + t.Fatalf("parse recorded template: %v", err) + } + + wantActions := map[headless.ActionType]bool{ + headless.ActionTextInput: false, headless.ActionKeyboard: false, + headless.ActionCheck: false, headless.ActionSelectInput: false, + headless.ActionHover: false, headless.ActionDblClick: false, + headless.ActionDispatchEvent: false, headless.ActionStorage: false, + headless.ActionCookie: false, headless.ActionSetViewport: false, + headless.ActionAssert: false, + } + for _, step := range tmpl.RequestsHeadless[0].Steps { + if _, tracked := wantActions[step.ActionType.ActionType]; tracked { + wantActions[step.ActionType.ActionType] = true + } + } + for actionType, found := range wantActions { + if !found { + t.Errorf("recorded template missing %s action", actionType) + } + } + + out := recExecString(t, cmd, ctx, []string{"template", templatePath, srv.URL}) + if !strings.Contains(out, "Template: extended-roundtrip") { + t.Fatalf("extended template did not replay: %s", out) + } +} + +func TestE2E_RecordReplayAuthenticatedDashboard(t *testing.T) { + skipIfNoBrowserR(t) + + type loginRequest struct { + Email string `json:"email"` + Password string `json:"password"` + Workspace string `json:"workspace"` + Remember bool `json:"remember"` + } + + var ( + loginMu sync.Mutex + loginRequests []loginRequest + ) + mux := http.NewServeMux() + mux.HandleFunc("/signin", func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprint(w, `Acme Cloud Sign In +
+

Sign in to Acme Cloud

+
+ + + + + + +

+
+
+`) + }) + mux.HandleFunc("/api/login", func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + var payload loginRequest + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + http.Error(w, "invalid JSON", http.StatusBadRequest) + return + } + loginMu.Lock() + loginRequests = append(loginRequests, payload) + loginMu.Unlock() + if payload.Email != "analyst@example.test" || payload.Password != "correct-horse" || payload.Workspace != "security" || !payload.Remember { + http.Error(w, "invalid credentials", http.StatusUnauthorized) + return + } + http.SetCookie(w, &http.Cookie{Name: "auth_session", Value: "session-e2e", Path: "/", HttpOnly: true, SameSite: http.SameSiteLaxMode}) + w.Header().Set("Content-Type", "application/json") + fmt.Fprint(w, `{"workspace":"security"}`) + }) + mux.HandleFunc("/app/dashboard", func(w http.ResponseWriter, r *http.Request) { + cookie, err := r.Cookie("auth_session") + if err != nil || cookie.Value != "session-e2e" { + http.Redirect(w, r, "/signin", http.StatusFound) + return + } + workspace := r.URL.Query().Get("workspace") + if workspace == "" { + workspace = "unknown" + } + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprintf(w, `Security Dashboard + +

%s dashboard

Welcome, analyst@example.test

+

Authenticated session

+`, strings.ToUpper(workspace[:1])+workspace[1:]) + }) + srv := httptest.NewServer(mux) + defer srv.Close() + + workDir := t.TempDir() + cmd := New(workDir) + defer cmd.Close() + ctx := context.Background() + + recExecString(t, cmd, ctx, []string{"open", srv.URL + "/signin", "--session", "auth", "--record", "--timeout", "10"}) + recExecString(t, cmd, ctx, []string{"fill", "auth", "label=Work email", "analyst@example.test"}) + recExecString(t, cmd, ctx, []string{"fill", "auth", "label=Password", "correct-horse"}) + recExecString(t, cmd, ctx, []string{"select-option", "auth", `role=combobox[name="Workspace"]`, "security"}) + recExecString(t, cmd, ctx, []string{"check", "auth", "testid=remember-device"}) + recExecString(t, cmd, ctx, []string{"click", "auth", `role=button[name="Sign in"]`}) + recExecString(t, cmd, ctx, []string{"wait-for-response", "auth", "/api/login"}) + recExecString(t, cmd, ctx, []string{"wait-for-url", "auth", "/app/dashboard"}) + welcome := recExecString(t, cmd, ctx, []string{"inner-text", "auth", "testid=welcome"}) + if !strings.Contains(welcome, "Welcome, analyst@example.test") { + t.Fatalf("dashboard did not render signed-in user: %s", welcome) + } + recExecString(t, cmd, ctx, []string{"is-visible", "auth", `role=heading[name="Security dashboard"]`}) + recExecString(t, cmd, ctx, []string{"is-visible", "auth", "testid=session-status"}) + recExecString(t, cmd, ctx, []string{"is-enabled", "auth", `role=button[name="Sign out"]`}) + storage := recExecString(t, cmd, ctx, []string{"localstorage-get", "auth", "lastWorkspace"}) + if !strings.Contains(storage, "security") { + t.Fatalf("dashboard storage state missing: %s", storage) + } + cookie := recExecString(t, cmd, ctx, []string{"cookie-get", "auth", "auth_session"}) + if !strings.Contains(cookie, "session-e2e") { + t.Fatalf("authenticated cookie missing: %s", cookie) + } + + templatePath := filepath.Join(workDir, "authenticated-dashboard.yaml") + recExecString(t, cmd, ctx, []string{"record", "auth", "--save", templatePath, "--id", "authenticated-dashboard-e2e"}) + data, err := os.ReadFile(templatePath) + if err != nil { + t.Fatal(err) + } + tmpl, err := headless.ParseTemplate(data) + if err != nil { + t.Fatalf("parse recorded template: %v", err) + } + if len(tmpl.RequestsHeadless) != 1 { + t.Fatalf("recorded template has %d headless requests", len(tmpl.RequestsHeadless)) + } + + wantActions := map[headless.ActionType]bool{ + headless.ActionNavigate: false, headless.ActionTextInput: false, + headless.ActionSelectInput: false, headless.ActionCheck: false, + headless.ActionClick: false, headless.ActionWaitResponse: false, + headless.ActionWaitURL: false, headless.ActionAssert: false, + headless.ActionExtract: false, + } + for _, step := range tmpl.RequestsHeadless[0].Steps { + if _, tracked := wantActions[step.ActionType.ActionType]; tracked { + wantActions[step.ActionType.ActionType] = true + } + if step.ActionType.ActionType == headless.ActionNavigate && strings.Contains(step.GetArg("url"), srv.URL) { + t.Fatalf("recorded navigation leaked the original origin: %s", step.GetArg("url")) + } + } + for actionType, found := range wantActions { + if !found { + t.Errorf("recorded authenticated flow missing %s action", actionType) + } + } + + out := recExecString(t, cmd, ctx, []string{"template", templatePath, srv.URL + "/signin"}) + if !strings.Contains(out, "Template: authenticated-dashboard-e2e") { + t.Fatalf("authenticated template did not replay: %s", out) + } + + loginMu.Lock() + requests := append([]loginRequest(nil), loginRequests...) + loginMu.Unlock() + if len(requests) != 2 { + t.Fatalf("login API received %d requests, want one live and one replay request", len(requests)) + } + for index, request := range requests { + if request.Email != "analyst@example.test" || request.Password != "correct-horse" || request.Workspace != "security" || !request.Remember { + t.Errorf("login request %d was not reproduced: %#v", index, request) + } + } +} diff --git a/tools/record/backend_native.go b/tools/record/backend_native.go new file mode 100644 index 00000000..784060b8 --- /dev/null +++ b/tools/record/backend_native.go @@ -0,0 +1,302 @@ +//go:build record_ffmpeg && cgo && (windows || linux) + +package record + +import ( + "context" + "errors" + "fmt" + "image" + "strconv" + "sync" + + "github.com/asticode/go-astiav" +) + +type nativeCaptureTarget struct { + format string + url string + options map[string]string +} + +type ffmpegBackend struct{} + +var registerDevicesOnce sync.Once + +func newPlatformBackend() captureBackend { + registerDevicesOnce.Do(func() { + astiav.RegisterAllDevices() + astiav.SetLogLevel(astiav.LogLevelWarning) + }) + return &ffmpegBackend{} +} + +func (b *ffmpegBackend) Resolve(ctx context.Context, req captureRequest) (resolvedTarget, error) { + return resolvePlatformTarget(ctx, req) +} + +func (b *ffmpegBackend) Screenshot(ctx context.Context, target resolvedTarget) (image.Image, error) { + pipeline, err := openInputPipeline(ctx, target, defaultFPS) + if err != nil { + return nil, err + } + defer pipeline.close() + var output image.Image + err = pipeline.nextFrame(ctx, func(frame *astiav.Frame) error { + sws, err := astiav.CreateSoftwareScaleContext( + frame.Width(), frame.Height(), frame.PixelFormat(), + frame.Width(), frame.Height(), astiav.PixelFormatRgba, + astiav.NewSoftwareScaleContextFlags(astiav.SoftwareScaleContextFlagBilinear), + ) + if err != nil { + return fmt.Errorf("create screenshot scaler: %w", err) + } + defer sws.Free() + dst := astiav.AllocFrame() + if dst == nil { + return fmt.Errorf("allocate screenshot frame") + } + defer dst.Free() + if err := sws.ScaleFrame(frame, dst); err != nil { + return fmt.Errorf("scale screenshot frame: %w", err) + } + img, err := dst.Data().GuessImageFormat() + if err != nil { + return fmt.Errorf("create screenshot image: %w", err) + } + if err := dst.Data().ToImage(img); err != nil { + return fmt.Errorf("copy screenshot pixels: %w", err) + } + output = img + return nil + }) + if err != nil { + return nil, err + } + return output, nil +} + +func (b *ffmpegBackend) Record(ctx context.Context, target resolvedTarget, output string, fps int) (result mediaInfo, retErr error) { + pipeline, err := openInputPipeline(ctx, target, fps) + if err != nil { + return result, err + } + defer pipeline.close() + + var encoder *videoEncoder + defer func() { + if encoder == nil { + return + } + if err := encoder.flush(); retErr == nil && err != nil { + retErr = err + } + if err := encoder.close(); retErr == nil && err != nil { + retErr = err + } + }() + + process := func(frame *astiav.Frame) error { + if encoder == nil { + width, height := frame.Width()&^1, frame.Height()&^1 + if width <= 0 || height <= 0 { + return fmt.Errorf("invalid capture dimensions %dx%d", frame.Width(), frame.Height()) + } + encoder, err = openVideoEncoder(output, fps, width, height, frame.PixelFormat()) + if err != nil { + return err + } + result.Width, result.Height = width, height + } + if frame.Width() < result.Width || frame.Height() < result.Height { + return fmt.Errorf("capture target shrank from %dx%d to %dx%d", result.Width, result.Height, frame.Width(), frame.Height()) + } + if err := encoder.write(frame, result.Frames); err != nil { + return err + } + result.Frames++ + return nil + } + + for { + err := pipeline.nextFrame(ctx, process) + if err == nil { + continue + } + if ctx.Err() != nil || errors.Is(err, astiav.ErrExit) { + break + } + if errors.Is(err, astiav.ErrEof) { + if result.Frames == 0 { + return result, fmt.Errorf("capture target ended before producing frames") + } + return result, fmt.Errorf("capture target closed or became unavailable") + } + return result, err + } + if result.Frames == 0 { + return result, fmt.Errorf("recording stopped before producing frames") + } + return result, nil +} + +type inputPipeline struct { + formatContext *astiav.FormatContext + decoder *astiav.CodecContext + stream *astiav.Stream + packet *astiav.Packet + frame *astiav.Frame + interrupter *astiav.IOInterrupter + wakeDone chan struct{} + watchStopped chan struct{} +} + +func openInputPipeline(ctx context.Context, target resolvedTarget, fps int) (*inputPipeline, error) { + native, ok := target.Native.(nativeCaptureTarget) + if !ok { + return nil, fmt.Errorf("invalid native capture target") + } + inputFormat := astiav.FindInputFormat(native.format) + if inputFormat == nil { + return nil, fmt.Errorf("FFmpeg input device %q is unavailable", native.format) + } + p := &inputPipeline{wakeDone: make(chan struct{}), watchStopped: make(chan struct{})} + p.formatContext = astiav.AllocFormatContext() + if p.formatContext == nil { + return nil, fmt.Errorf("allocate input format context") + } + p.interrupter = astiav.NewIOInterrupter() + p.formatContext.SetIOInterrupter(p.interrupter) + go func() { + defer close(p.watchStopped) + select { + case <-ctx.Done(): + p.interrupter.Interrupt() + case <-p.wakeDone: + } + }() + opts := astiav.NewDictionary() + defer opts.Free() + if err := opts.Set("framerate", strconv.Itoa(fps), astiav.NewDictionaryFlags()); err != nil { + p.close() + return nil, fmt.Errorf("set capture framerate: %w", err) + } + if err := opts.Set("draw_mouse", "1", astiav.NewDictionaryFlags()); err != nil { + p.close() + return nil, fmt.Errorf("set mouse capture option: %w", err) + } + for key, value := range native.options { + if err := opts.Set(key, value, astiav.NewDictionaryFlags()); err != nil { + p.close() + return nil, fmt.Errorf("set capture option %s: %w", key, err) + } + } + if err := p.formatContext.OpenInput(native.url, inputFormat, opts); err != nil { + p.close() + return nil, fmt.Errorf("open %s capture input: %w", native.format, err) + } + for _, stream := range p.formatContext.Streams() { + if stream.CodecParameters().MediaType() == astiav.MediaTypeVideo { + p.stream = stream + break + } + } + if p.stream == nil { + p.close() + return nil, fmt.Errorf("capture input has no video stream") + } + codec := astiav.FindDecoder(p.stream.CodecParameters().CodecID()) + if codec == nil { + p.close() + return nil, fmt.Errorf("decoder for %s is unavailable", p.stream.CodecParameters().CodecID()) + } + p.decoder = astiav.AllocCodecContext(codec) + if p.decoder == nil { + p.close() + return nil, fmt.Errorf("allocate capture decoder") + } + if err := p.stream.CodecParameters().ToCodecContext(p.decoder); err != nil { + p.close() + return nil, fmt.Errorf("configure capture decoder: %w", err) + } + p.decoder.SetTimeBase(p.stream.TimeBase()) + p.decoder.SetFramerate(astiav.NewRational(fps, 1)) + if err := p.decoder.Open(codec, nil); err != nil { + p.close() + return nil, fmt.Errorf("open capture decoder: %w", err) + } + p.packet = astiav.AllocPacket() + p.frame = astiav.AllocFrame() + if p.packet == nil || p.frame == nil { + p.close() + return nil, fmt.Errorf("allocate capture packet/frame") + } + return p, nil +} + +func (p *inputPipeline) nextFrame(ctx context.Context, consume func(*astiav.Frame) error) error { + for { + if err := ctx.Err(); err != nil { + p.interrupter.Interrupt() + return err + } + if err := p.formatContext.ReadFrame(p.packet); err != nil { + if ctx.Err() != nil { + return ctx.Err() + } + return err + } + if p.packet.StreamIndex() != p.stream.Index() { + p.packet.Unref() + continue + } + p.packet.RescaleTs(p.stream.TimeBase(), p.decoder.TimeBase()) + err := p.decoder.SendPacket(p.packet) + p.packet.Unref() + if err != nil { + return fmt.Errorf("send capture packet: %w", err) + } + for { + err = p.decoder.ReceiveFrame(p.frame) + if errors.Is(err, astiav.ErrEagain) { + break + } + if err != nil { + return err + } + err = consume(p.frame) + p.frame.Unref() + return err + } + } +} + +func (p *inputPipeline) close() { + if p == nil { + return + } + select { + case <-p.wakeDone: + default: + close(p.wakeDone) + } + if p.watchStopped != nil { + <-p.watchStopped + } + if p.frame != nil { + p.frame.Free() + } + if p.packet != nil { + p.packet.Free() + } + if p.decoder != nil { + p.decoder.Free() + } + if p.formatContext != nil { + p.formatContext.CloseInput() + p.formatContext.Free() + } + if p.interrupter != nil { + p.interrupter.Free() + } +} diff --git a/tools/record/backend_native_integration_test.go b/tools/record/backend_native_integration_test.go new file mode 100644 index 00000000..3fde73a2 --- /dev/null +++ b/tools/record/backend_native_integration_test.go @@ -0,0 +1,82 @@ +//go:build record_ffmpeg && record_integration && cgo && (windows || linux) + +package record + +import ( + "bytes" + "context" + "encoding/binary" + "os" + "path/filepath" + "testing" + "time" +) + +func TestNativeDesktopScreenshotAndRecord(t *testing.T) { + backend := newPlatformBackend() + target, err := backend.Resolve(context.Background(), captureRequest{Target: "desktop", FPS: 10}) + if err != nil { + t.Fatal(err) + } + img, err := backend.Screenshot(context.Background(), target) + if err != nil { + t.Fatal(err) + } + if img.Bounds().Dx() <= 0 || img.Bounds().Dy() <= 0 { + t.Fatalf("invalid screenshot bounds %v", img.Bounds()) + } + + output := filepath.Join(t.TempDir(), "desktop.mp4") + ctx, cancel := context.WithTimeout(context.Background(), 1500*time.Millisecond) + defer cancel() + media, err := backend.Record(ctx, target, output, 10) + if err != nil { + t.Fatal(err) + } + if media.Frames == 0 { + t.Fatal("recording produced no frames") + } + if info, err := os.Stat(output); err != nil || info.Size() == 0 { + t.Fatalf("recording output info=%v err=%v", info, err) + } + + assertMP4H264(t, output) +} + +func assertMP4H264(t *testing.T, path string) { + t.Helper() + data, err := os.ReadFile(path) + if err != nil { + t.Fatal(err) + } + topLevel := make(map[string]bool) + for offset := 0; offset < len(data); { + if len(data)-offset < 8 { + t.Fatalf("truncated MP4 atom header at byte %d", offset) + } + size := uint64(binary.BigEndian.Uint32(data[offset : offset+4])) + headerSize := uint64(8) + if size == 1 { + if len(data)-offset < 16 { + t.Fatalf("truncated extended MP4 atom header at byte %d", offset) + } + size = binary.BigEndian.Uint64(data[offset+8 : offset+16]) + headerSize = 16 + } else if size == 0 { + size = uint64(len(data) - offset) + } + if size < headerSize || size > uint64(len(data)-offset) { + t.Fatalf("invalid MP4 atom size %d at byte %d", size, offset) + } + topLevel[string(data[offset+4:offset+8])] = true + offset += int(size) + } + for _, atom := range []string{"ftyp", "moov", "mdat"} { + if !topLevel[atom] { + t.Fatalf("recorded MP4 is missing %s atom", atom) + } + } + if !bytes.Contains(data, []byte("avc1")) { + t.Fatal("recorded MP4 has no H.264 avc1 sample entry") + } +} diff --git a/tools/record/backend_unavailable.go b/tools/record/backend_unavailable.go new file mode 100644 index 00000000..fe67b506 --- /dev/null +++ b/tools/record/backend_unavailable.go @@ -0,0 +1,25 @@ +//go:build !record_ffmpeg || !cgo || (!windows && !linux) + +package record + +import ( + "context" + "fmt" + "image" +) + +type unavailableBackend struct{} + +func newPlatformBackend() captureBackend { return unavailableBackend{} } + +func (unavailableBackend) Resolve(context.Context, captureRequest) (resolvedTarget, error) { + return resolvedTarget{}, fmt.Errorf("native recorder is not linked; build the full edition with CGO and the record_ffmpeg tag") +} + +func (unavailableBackend) Screenshot(context.Context, resolvedTarget) (image.Image, error) { + return nil, fmt.Errorf("native recorder is unavailable") +} + +func (unavailableBackend) Record(context.Context, resolvedTarget, string, int) (mediaInfo, error) { + return mediaInfo{}, fmt.Errorf("native recorder is unavailable") +} diff --git a/tools/record/config.go b/tools/record/config.go new file mode 100644 index 00000000..9ddca429 --- /dev/null +++ b/tools/record/config.go @@ -0,0 +1,33 @@ +package record + +import ( + "fmt" + "strconv" + "strings" +) + +const ( + maxConcurrentEnv = "AISCAN_RECORD_MAX_CONCURRENT" + defaultMaxConcurrent = 4 + maxConcurrentLimit = 16 +) + +type environmentLookup func(string) (string, bool) + +func maxConcurrentFromEnvironment(lookup environmentLookup) (int, error) { + if lookup == nil { + return defaultMaxConcurrent, nil + } + raw, ok := lookup(maxConcurrentEnv) + if !ok || strings.TrimSpace(raw) == "" { + return defaultMaxConcurrent, nil + } + value, err := strconv.Atoi(strings.TrimSpace(raw)) + if err != nil { + return 0, fmt.Errorf("%s must be an integer between 1 and %d", maxConcurrentEnv, maxConcurrentLimit) + } + if value < 1 || value > maxConcurrentLimit { + return 0, fmt.Errorf("%s must be between 1 and %d", maxConcurrentEnv, maxConcurrentLimit) + } + return value, nil +} diff --git a/tools/record/config_test.go b/tools/record/config_test.go new file mode 100644 index 00000000..9fd8dfc2 --- /dev/null +++ b/tools/record/config_test.go @@ -0,0 +1,45 @@ +package record + +import ( + "strings" + "testing" +) + +func TestMaxConcurrentFromEnvironment(t *testing.T) { + tests := []struct { + name string + value string + present bool + want int + wantErr string + }{ + {name: "unset", want: defaultMaxConcurrent}, + {name: "blank", value: " ", present: true, want: defaultMaxConcurrent}, + {name: "configured", value: " 6 ", present: true, want: 6}, + {name: "not integer", value: "many", present: true, wantErr: "must be an integer"}, + {name: "zero", value: "0", present: true, wantErr: "must be between"}, + {name: "above limit", value: "17", present: true, wantErr: "must be between"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := maxConcurrentFromEnvironment(func(name string) (string, bool) { + if name != maxConcurrentEnv || !tt.present { + return "", false + } + return tt.value, true + }) + if tt.wantErr == "" { + if err != nil { + t.Fatal(err) + } + if got != tt.want { + t.Fatalf("max concurrent = %d, want %d", got, tt.want) + } + return + } + if err == nil || !strings.Contains(err.Error(), tt.wantErr) { + t.Fatalf("error = %v, want substring %q", err, tt.wantErr) + } + }) + } +} diff --git a/tools/record/encoder_native.go b/tools/record/encoder_native.go new file mode 100644 index 00000000..60be19e5 --- /dev/null +++ b/tools/record/encoder_native.go @@ -0,0 +1,199 @@ +//go:build record_ffmpeg && cgo && (windows || linux) + +package record + +import ( + "errors" + "fmt" + + "github.com/asticode/go-astiav" +) + +type videoEncoder struct { + formatContext *astiav.FormatContext + codecContext *astiav.CodecContext + stream *astiav.Stream + ioContext *astiav.IOContext + scaler *astiav.SoftwareScaleContext + frame *astiav.Frame + packet *astiav.Packet + headerWritten bool + trailerWritten bool +} + +func openVideoEncoder(path string, fps, width, height int, sourceFormat astiav.PixelFormat) (*videoEncoder, error) { + encoder := &videoEncoder{} + var err error + encoder.formatContext, err = astiav.AllocOutputFormatContext(nil, "mp4", path) + if err != nil { + return nil, fmt.Errorf("allocate MP4 output: %w", err) + } + if encoder.formatContext == nil { + return nil, fmt.Errorf("allocate MP4 output") + } + codec := astiav.FindEncoderByName("libx264") + if codec == nil { + encoder.close() + return nil, fmt.Errorf("FFmpeg encoder libx264 is unavailable") + } + encoder.codecContext = astiav.AllocCodecContext(codec) + if encoder.codecContext == nil { + encoder.close() + return nil, fmt.Errorf("allocate libx264 context") + } + encoder.codecContext.SetWidth(width) + encoder.codecContext.SetHeight(height) + encoder.codecContext.SetPixelFormat(astiav.PixelFormatYuv420P) + encoder.codecContext.SetTimeBase(astiav.NewRational(1, fps)) + encoder.codecContext.SetFramerate(astiav.NewRational(fps, 1)) + encoder.codecContext.SetGopSize(fps * 2) + encoder.codecContext.SetMaxBFrames(2) + if encoder.formatContext.OutputFormat().Flags().Has(astiav.IOFormatFlagGlobalheader) { + encoder.codecContext.SetFlags(encoder.codecContext.Flags().Add(astiav.CodecContextFlagGlobalHeader)) + } + codecOpts := astiav.NewDictionary() + if err := codecOpts.Set("preset", "veryfast", astiav.NewDictionaryFlags()); err != nil { + codecOpts.Free() + encoder.close() + return nil, fmt.Errorf("set encoder preset: %w", err) + } + if err := codecOpts.Set("crf", "23", astiav.NewDictionaryFlags()); err != nil { + codecOpts.Free() + encoder.close() + return nil, fmt.Errorf("set encoder quality: %w", err) + } + if err := encoder.codecContext.Open(codec, codecOpts); err != nil { + codecOpts.Free() + encoder.close() + return nil, fmt.Errorf("open libx264 encoder: %w", err) + } + codecOpts.Free() + encoder.stream = encoder.formatContext.NewStream(nil) + if encoder.stream == nil { + encoder.close() + return nil, fmt.Errorf("create MP4 video stream") + } + if err := encoder.stream.CodecParameters().FromCodecContext(encoder.codecContext); err != nil { + encoder.close() + return nil, fmt.Errorf("copy encoder parameters: %w", err) + } + encoder.stream.SetTimeBase(encoder.codecContext.TimeBase()) + encoder.ioContext, err = astiav.OpenIOContext(path, astiav.NewIOContextFlags(astiav.IOContextFlagWrite), nil, nil) + if err != nil { + encoder.close() + return nil, fmt.Errorf("open MP4 output: %w", err) + } + encoder.formatContext.SetPb(encoder.ioContext) + headerOpts := astiav.NewDictionary() + if err := headerOpts.Set("movflags", "+faststart", astiav.NewDictionaryFlags()); err != nil { + headerOpts.Free() + encoder.close() + return nil, fmt.Errorf("set MP4 output options: %w", err) + } + if err := encoder.formatContext.WriteHeader(headerOpts); err != nil { + headerOpts.Free() + encoder.close() + return nil, fmt.Errorf("write MP4 header: %w", err) + } + headerOpts.Free() + encoder.headerWritten = true + encoder.scaler, err = astiav.CreateSoftwareScaleContext( + width, height, sourceFormat, + width, height, astiav.PixelFormatYuv420P, + astiav.NewSoftwareScaleContextFlags(astiav.SoftwareScaleContextFlagBilinear), + ) + if err != nil { + encoder.close() + return nil, fmt.Errorf("create video scaler: %w", err) + } + encoder.frame = astiav.AllocFrame() + encoder.packet = astiav.AllocPacket() + if encoder.frame == nil || encoder.packet == nil { + encoder.close() + return nil, fmt.Errorf("allocate encoder frame/packet") + } + encoder.frame.SetWidth(width) + encoder.frame.SetHeight(height) + encoder.frame.SetPixelFormat(astiav.PixelFormatYuv420P) + if err := encoder.frame.AllocBuffer(32); err != nil { + encoder.close() + return nil, fmt.Errorf("allocate encoder frame buffer: %w", err) + } + return encoder, nil +} + +func (encoder *videoEncoder) write(source *astiav.Frame, pts int64) error { + if err := encoder.frame.MakeWritable(); err != nil { + return fmt.Errorf("make encoder frame writable: %w", err) + } + if err := encoder.scaler.ScaleFrame(source, encoder.frame); err != nil { + return fmt.Errorf("convert capture frame: %w", err) + } + encoder.frame.SetPts(pts) + encoder.frame.SetPictureType(astiav.PictureTypeNone) + return encoder.sendFrame(encoder.frame) +} + +func (encoder *videoEncoder) sendFrame(frame *astiav.Frame) error { + if err := encoder.codecContext.SendFrame(frame); err != nil { + return fmt.Errorf("send encoder frame: %w", err) + } + for { + err := encoder.codecContext.ReceivePacket(encoder.packet) + if errors.Is(err, astiav.ErrEagain) || errors.Is(err, astiav.ErrEof) { + return nil + } + if err != nil { + return fmt.Errorf("receive encoded packet: %w", err) + } + encoder.packet.SetStreamIndex(encoder.stream.Index()) + encoder.packet.RescaleTs(encoder.codecContext.TimeBase(), encoder.stream.TimeBase()) + err = encoder.formatContext.WriteInterleavedFrame(encoder.packet) + encoder.packet.Unref() + if err != nil { + return fmt.Errorf("write encoded packet: %w", err) + } + } +} + +func (encoder *videoEncoder) flush() error { + if encoder == nil || encoder.codecContext == nil { + return nil + } + return encoder.sendFrame(nil) +} + +func (encoder *videoEncoder) close() error { + if encoder == nil { + return nil + } + var result error + if encoder.headerWritten && !encoder.trailerWritten && encoder.formatContext != nil { + if err := encoder.formatContext.WriteTrailer(); err != nil { + result = fmt.Errorf("write MP4 trailer: %w", err) + } else { + encoder.trailerWritten = true + } + } + if encoder.packet != nil { + encoder.packet.Free() + } + if encoder.frame != nil { + encoder.frame.Free() + } + if encoder.scaler != nil { + encoder.scaler.Free() + } + if encoder.codecContext != nil { + encoder.codecContext.Free() + } + if encoder.ioContext != nil { + if err := encoder.ioContext.Close(); result == nil && err != nil { + result = err + } + } + if encoder.formatContext != nil { + encoder.formatContext.Free() + } + return result +} diff --git a/tools/record/helpers.go b/tools/record/helpers.go new file mode 100644 index 00000000..db2413d2 --- /dev/null +++ b/tools/record/helpers.go @@ -0,0 +1,147 @@ +package record + +import ( + "context" + "crypto/rand" + "encoding/hex" + "encoding/json" + "fmt" + "math" + "path/filepath" + "runtime" + "strconv" + "strings" + "time" + + "github.com/chainreactors/aiscan/core/tool" +) + +func normalizeDuration(seconds float64, required bool) (time.Duration, error) { + if math.IsNaN(seconds) || math.IsInf(seconds, 0) { + return 0, fmt.Errorf("duration_seconds must be a finite number") + } + if seconds < 0 { + return 0, fmt.Errorf("duration_seconds cannot be negative") + } + if required && seconds == 0 { + return 0, fmt.Errorf("duration_seconds must be greater than zero for action=record") + } + if seconds > float64(math.MaxInt64)/float64(time.Second) { + return 0, fmt.Errorf("duration_seconds is too large") + } + duration := time.Duration(seconds * float64(time.Second)) + if seconds > 0 && duration <= 0 { + return 0, fmt.Errorf("duration_seconds is too small") + } + return duration, nil +} + +func normalizeCaptureArgs(args Args) (captureRequest, error) { + if args.PID > int64(^uint32(0)) { + return captureRequest{}, fmt.Errorf("pid exceeds the supported 32-bit process identifier range") + } + target := strings.ToLower(strings.TrimSpace(args.Target)) + if target == "" { + target = "desktop" + } + fps := args.FPS + if fps == 0 { + fps = defaultFPS + } + if fps < 1 || fps > 60 { + return captureRequest{}, fmt.Errorf("fps must be between 1 and 60") + } + req := captureRequest{Target: target, PID: args.PID, FPS: fps} + if strings.TrimSpace(args.WindowHandle) != "" { + value, err := strconv.ParseUint(strings.TrimSpace(args.WindowHandle), 0, 64) + if err != nil || value == 0 { + return captureRequest{}, fmt.Errorf("invalid window_handle %q", args.WindowHandle) + } + req.WindowHandle = value + } + switch target { + case "desktop": + if req.WindowHandle != 0 || req.PID != 0 { + return captureRequest{}, fmt.Errorf("window_handle and pid require target=window") + } + case "window": + if req.WindowHandle != 0 && req.PID != 0 { + return captureRequest{}, fmt.Errorf("window_handle and pid are mutually exclusive") + } + if req.WindowHandle == 0 && req.PID <= 0 { + return captureRequest{}, fmt.Errorf("target=window requires window_handle or pid") + } + default: + return captureRequest{}, fmt.Errorf("unsupported target %q", args.Target) + } + return req, nil +} + +func (t *Tool) outputPath(ctx context.Context, requested, base, ext string) (string, error) { + path := strings.TrimSpace(requested) + if path == "" { + if invocationDir := tool.InvocationFromContext(ctx).WorkDir; invocationDir != "" { + path = filepath.Join(invocationDir, ".aiscan", "record", base+ext) + } else { + path = filepath.Join(t.outputDir, base+ext) + } + } else { + if filepath.Ext(path) == "" { + path += ext + } else if !strings.EqualFold(filepath.Ext(path), ext) { + return "", fmt.Errorf("output must use %s extension", ext) + } + if !filepath.IsAbs(path) { + path = filepath.Join(tool.WorkDirFromContext(ctx, t.workDir), path) + } + } + abs, err := filepath.Abs(filepath.Clean(path)) + if err != nil { + return "", fmt.Errorf("resolve output path: %w", err) + } + return abs, nil +} + +func (t *Tool) mediaURI(ctx context.Context, path string) string { + base := tool.WorkDirFromContext(ctx, t.workDir) + if base != "" { + if relative, err := filepath.Rel(base, path); err == nil && relative != ".." && !strings.HasPrefix(relative, ".."+string(filepath.Separator)) { + return filepath.ToSlash(relative) + } + } + return filepath.ToSlash(path) +} + +func newID() string { + buf := make([]byte, 8) + if _, err := rand.Read(buf); err == nil { + return hex.EncodeToString(buf) + } + return strconv.FormatInt(time.Now().UnixNano(), 36) +} + +func jsonResult(value any) (*tool.Result, error) { + return tool.TextResult(marshalJSON(value)), nil +} + +func marshalJSON(value any) string { + data, err := json.MarshalIndent(value, "", " ") + if err != nil { + return fmt.Sprintf(`{"error":%q}`, err.Error()) + } + return string(data) +} + +func samePath(a, b string) bool { + aAbs, aErr := filepath.Abs(a) + bAbs, bErr := filepath.Abs(b) + if aErr == nil && bErr == nil { + a, b = filepath.Clean(aAbs), filepath.Clean(bAbs) + } else { + a, b = filepath.Clean(a), filepath.Clean(b) + } + if runtime.GOOS == "windows" { + return strings.EqualFold(a, b) + } + return a == b +} diff --git a/tools/record/register.go b/tools/record/register.go new file mode 100644 index 00000000..12213f95 --- /dev/null +++ b/tools/record/register.go @@ -0,0 +1,34 @@ +//go:build full && (windows || linux) + +package record + +import ( + "os" + + "github.com/chainreactors/aiscan/core/capability" + coreconfig "github.com/chainreactors/aiscan/core/config" + "github.com/chainreactors/aiscan/pkg/commands" +) + +func init() { + capability.Register(capability.Descriptor{ + ID: "record", Kind: capability.KindTool, Group: "record", + Optional: true, Default: true, + }) + commands.RegisterFactory(commands.Factory{ + Capability: "record", + Build: func(deps *commands.Deps, reg *commands.CommandRegistry) { + maxConcurrent, err := maxConcurrentFromEnvironment(os.LookupEnv) + if err != nil { + deps.GetLogger().Warnf("record config: %s; using default %d", err, defaultMaxConcurrent) + maxConcurrent = defaultMaxConcurrent + } + reg.RegisterTool(New( + deps.WorkDir, + coreconfig.DataSubDir("record"), + maxConcurrent, + newPlatformBackend(), + )) + }, + }) +} diff --git a/tools/record/register_test.go b/tools/record/register_test.go new file mode 100644 index 00000000..90a65a17 --- /dev/null +++ b/tools/record/register_test.go @@ -0,0 +1,40 @@ +//go:build full && (windows || linux) + +package record + +import ( + "testing" + + "github.com/chainreactors/aiscan/core/capability" + "github.com/chainreactors/aiscan/pkg/commands" +) + +func TestRegister(t *testing.T) { + t.Setenv(maxConcurrentEnv, "3") + reg := commands.NewRegistry() + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"record"}}), &commands.Deps{ + WorkDir: t.TempDir(), + }, reg) + tool, ok := reg.GetTool("record") + if !ok { + t.Fatal("record tool is not registered") + } + if got := tool.(*Tool).maxConcurrent; got != 3 { + t.Fatalf("max concurrent = %d, want 3", got) + } +} + +func TestRegisterInvalidMaxConcurrentUsesDefault(t *testing.T) { + t.Setenv(maxConcurrentEnv, "17") + reg := commands.NewRegistry() + commands.BuildPlan(capability.Select(capability.Options{Groups: []string{"record"}}), &commands.Deps{ + WorkDir: t.TempDir(), + }, reg) + tool, ok := reg.GetTool("record") + if !ok { + t.Fatal("record tool is not registered") + } + if got := tool.(*Tool).maxConcurrent; got != defaultMaxConcurrent { + t.Fatalf("max concurrent = %d, want default %d", got, defaultMaxConcurrent) + } +} diff --git a/tools/record/session.go b/tools/record/session.go new file mode 100644 index 00000000..7042ce83 --- /dev/null +++ b/tools/record/session.go @@ -0,0 +1,331 @@ +package record + +import ( + "context" + "fmt" + "image" + "os" + "path/filepath" + "sort" + "strings" + "sync" + "time" + + aop "github.com/chainreactors/aiscan/aop" + "github.com/chainreactors/aiscan/core/tool" +) + +type recordingSession struct { + mu sync.RWMutex + info SessionInfo + cancel context.CancelFunc + done chan struct{} +} + +func (t *Tool) start(ctx context.Context, args Args, duration time.Duration) (*recordingSession, error) { + if t.backend == nil { + return nil, fmt.Errorf("capture backend is unavailable") + } + req, err := normalizeCaptureArgs(args) + if err != nil { + return nil, err + } + target, err := callBackendResolve(t.backend, ctx, req) + if err != nil { + return nil, fmt.Errorf("resolve capture target: %w", err) + } + id := newID() + path, err := t.outputPath(ctx, args.Output, "record-"+id, ".mp4") + if err != nil { + return nil, err + } + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return nil, fmt.Errorf("create recording directory: %w", err) + } + + t.mu.Lock() + if t.closed { + t.mu.Unlock() + return nil, fmt.Errorf("record tool is closed") + } + if t.activeCountLocked() >= t.maxConcurrent { + ids := t.activeIDsLocked() + t.mu.Unlock() + return nil, fmt.Errorf("recording concurrency limit %d reached (active: %s)", t.maxConcurrent, strings.Join(ids, ", ")) + } + for _, existing := range t.sessions { + info := existing.snapshot() + if isActive(info.State) && samePath(info.Output, path) { + t.mu.Unlock() + return nil, fmt.Errorf("output path is already used by recording %s", info.RecordingID) + } + } + + runCtx, cancel := context.WithCancel(context.Background()) + if duration > 0 { + runCtx, cancel = context.WithTimeout(context.Background(), duration) + } + session := &recordingSession{ + info: SessionInfo{ + RecordingID: id, + State: sessionStarting, + Target: target.Info, + Output: path, + FPS: req.FPS, + }, + cancel: cancel, + done: make(chan struct{}), + } + t.sessions[id] = session + t.mu.Unlock() + + go t.runSession(runCtx, session, target) + return session, nil +} + +func (t *Tool) runSession(ctx context.Context, session *recordingSession, target resolvedTarget) { + defer session.cancel() + started := time.Now().UTC() + session.update(func(info *SessionInfo) { + if info.State == sessionStarting { + info.State = sessionRecording + } + info.StartedAt = &started + }) + snapshot := session.snapshot() + media, recordErr := callBackendRecord(t.backend, ctx, target, snapshot.Output, snapshot.FPS) + ended := time.Now().UTC() + var size int64 + if stat, err := os.Stat(snapshot.Output); err == nil { + size = stat.Size() + if recordErr == nil && size == 0 { + recordErr = fmt.Errorf("recording output is empty") + } + } else if recordErr == nil { + recordErr = fmt.Errorf("inspect recording output: %w", err) + } + session.update(func(info *SessionInfo) { + info.EndedAt = &ended + info.Bytes = size + info.Frames = media.Frames + if media.Width > 0 { + info.Target.Width = media.Width + } + if media.Height > 0 { + info.Target.Height = media.Height + } + if info.StartedAt != nil { + info.DurationMS = ended.Sub(*info.StartedAt).Milliseconds() + } + if recordErr != nil { + info.State = sessionFailed + info.Error = recordErr.Error() + } else { + info.State = sessionCompleted + } + }) + t.pruneSessionHistory() + close(session.done) +} + +func callBackendRecord(backend captureBackend, ctx context.Context, target resolvedTarget, output string, fps int) (media mediaInfo, err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = fmt.Errorf("capture backend panicked: %v", recovered) + } + }() + return backend.Record(ctx, target, output, fps) +} + +func callBackendResolve(backend captureBackend, ctx context.Context, req captureRequest) (target resolvedTarget, err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = fmt.Errorf("capture backend panicked while resolving target: %v", recovered) + } + }() + return backend.Resolve(ctx, req) +} + +func callBackendScreenshot(backend captureBackend, ctx context.Context, target resolvedTarget) (img image.Image, err error) { + defer func() { + if recovered := recover(); recovered != nil { + err = fmt.Errorf("capture backend panicked while taking screenshot: %v", recovered) + } + }() + return backend.Screenshot(ctx, target) +} + +func (t *Tool) stop(ctx context.Context, id string) (*tool.Result, error) { + if id == "" { + return nil, fmt.Errorf("recording_id is required for action=stop") + } + session, ok := t.session(id) + if !ok { + return nil, fmt.Errorf("recording %q not found", id) + } + if isActive(session.snapshot().State) { + session.update(func(info *SessionInfo) { info.State = sessionStopping }) + session.cancel() + } + return t.waitResult(ctx, session) +} + +func (t *Tool) waitResult(ctx context.Context, session *recordingSession) (*tool.Result, error) { + select { + case <-session.done: + info := session.snapshot() + if info.State == sessionFailed { + return tool.ErrorResult(marshalJSON(info)), nil + } + text := marshalJSON(info) + return &tool.Result{Output: []*aop.Content{ + aop.Text(text), + aop.MediaURI("video", "video/mp4", filepath.Base(info.Output), t.mediaURI(ctx, info.Output)), + }}, nil + case <-ctx.Done(): + session.cancel() + timer := time.NewTimer(sessionStopTimeout) + defer timer.Stop() + select { + case <-session.done: + return nil, ctx.Err() + case <-timer.C: + return nil, fmt.Errorf("stop recording after context cancellation: %w", ctx.Err()) + } + } +} + +func (t *Tool) status(id string) (*tool.Result, error) { + if id != "" { + session, ok := t.session(id) + if !ok { + return nil, fmt.Errorf("recording %q not found", id) + } + return jsonResult(session.snapshot()) + } + t.mu.RLock() + infos := make([]SessionInfo, 0, len(t.sessions)) + for _, session := range t.sessions { + infos = append(infos, session.snapshot()) + } + t.mu.RUnlock() + sort.Slice(infos, func(i, j int) bool { return infos[i].RecordingID < infos[j].RecordingID }) + return jsonResult(infos) +} + +func (t *Tool) Close() { + t.mu.Lock() + if t.closed { + t.mu.Unlock() + return + } + t.closed = true + sessions := make([]*recordingSession, 0, len(t.sessions)) + for _, session := range t.sessions { + if isActive(session.snapshot().State) { + session.update(func(info *SessionInfo) { info.State = sessionStopping }) + session.cancel() + sessions = append(sessions, session) + } + } + t.mu.Unlock() + + if len(sessions) == 0 { + return + } + timer := time.NewTimer(sessionStopTimeout) + defer timer.Stop() + for _, session := range sessions { + select { + case <-session.done: + case <-timer.C: + return + } + } +} + +func (t *Tool) activeCountLocked() int { + count := 0 + for _, session := range t.sessions { + if isActive(session.snapshot().State) { + count++ + } + } + return count +} + +func (t *Tool) activeIDsLocked() []string { + ids := make([]string, 0, t.maxConcurrent) + for id, session := range t.sessions { + if isActive(session.snapshot().State) { + ids = append(ids, id) + } + } + sort.Strings(ids) + return ids +} + +func (t *Tool) session(id string) (*recordingSession, bool) { + t.mu.RLock() + defer t.mu.RUnlock() + session, ok := t.sessions[id] + return session, ok +} + +func (t *Tool) pruneSessionHistory() { + t.mu.Lock() + defer t.mu.Unlock() + if len(t.sessions) <= maxSessionHistory { + return + } + type candidate struct { + id string + endedAt time.Time + } + candidates := make([]candidate, 0, len(t.sessions)) + for id, session := range t.sessions { + info := session.snapshot() + if isActive(info.State) { + continue + } + endedAt := time.Time{} + if info.EndedAt != nil { + endedAt = *info.EndedAt + } + candidates = append(candidates, candidate{id: id, endedAt: endedAt}) + } + sort.Slice(candidates, func(i, j int) bool { + if candidates[i].endedAt.Equal(candidates[j].endedAt) { + return candidates[i].id < candidates[j].id + } + return candidates[i].endedAt.Before(candidates[j].endedAt) + }) + remove := len(t.sessions) - maxSessionHistory + if remove > len(candidates) { + remove = len(candidates) + } + for _, candidate := range candidates[:remove] { + delete(t.sessions, candidate.id) + } +} + +func (session *recordingSession) snapshot() SessionInfo { + session.mu.RLock() + defer session.mu.RUnlock() + return session.info +} + +func (session *recordingSession) update(fn func(*SessionInfo)) { + session.mu.Lock() + defer session.mu.Unlock() + fn(&session.info) +} + +func isActive(state string) bool { + switch state { + case sessionStarting, sessionRecording, sessionStopping: + return true + default: + return false + } +} diff --git a/tools/record/target_linux.go b/tools/record/target_linux.go new file mode 100644 index 00000000..97bfd2fb --- /dev/null +++ b/tools/record/target_linux.go @@ -0,0 +1,199 @@ +//go:build record_ffmpeg && cgo && linux + +package record + +/* +#cgo pkg-config: xcb +#include +#include +#include +#include + +typedef struct { + uint32_t window; + uint32_t pid; + uint32_t width; + uint32_t height; + char title[512]; +} aiscan_x11_window; + +static xcb_atom_t aiscan_atom(xcb_connection_t *c, const char *name) { + xcb_intern_atom_cookie_t cookie = xcb_intern_atom(c, 0, strlen(name), name); + xcb_intern_atom_reply_t *reply = xcb_intern_atom_reply(c, cookie, NULL); + if (!reply) return XCB_ATOM_NONE; + xcb_atom_t atom = reply->atom; + free(reply); + return atom; +} + +static xcb_screen_t *aiscan_screen(xcb_connection_t *c, int number) { + const xcb_setup_t *setup = xcb_get_setup(c); + xcb_screen_iterator_t it = xcb_setup_roots_iterator(setup); + for (int i = 0; i < number && it.rem; i++) xcb_screen_next(&it); + return it.rem ? it.data : NULL; +} + +static int aiscan_window_info(xcb_connection_t *c, xcb_window_t window, + xcb_atom_t pid_atom, xcb_atom_t name_atom, + aiscan_x11_window *out) { + xcb_get_window_attributes_reply_t *attrs = xcb_get_window_attributes_reply( + c, xcb_get_window_attributes(c, window), NULL); + if (!attrs || attrs->map_state != XCB_MAP_STATE_VIEWABLE) { + free(attrs); + return 0; + } + free(attrs); + xcb_get_geometry_reply_t *geometry = xcb_get_geometry_reply(c, xcb_get_geometry(c, window), NULL); + if (!geometry || geometry->width == 0 || geometry->height == 0) { + free(geometry); + return 0; + } + memset(out, 0, sizeof(*out)); + out->window = window; + out->width = geometry->width; + out->height = geometry->height; + free(geometry); + + if (pid_atom != XCB_ATOM_NONE) { + xcb_get_property_reply_t *pid_reply = xcb_get_property_reply(c, + xcb_get_property(c, 0, window, pid_atom, XCB_ATOM_CARDINAL, 0, 1), NULL); + if (pid_reply && xcb_get_property_value_length(pid_reply) >= 4) { + out->pid = *(uint32_t *)xcb_get_property_value(pid_reply); + } + free(pid_reply); + } + if (name_atom != XCB_ATOM_NONE) { + xcb_get_property_reply_t *name_reply = xcb_get_property_reply(c, + xcb_get_property(c, 0, window, name_atom, XCB_GET_PROPERTY_TYPE_ANY, 0, 511), NULL); + if (name_reply) { + int length = xcb_get_property_value_length(name_reply); + if (length > 511) length = 511; + if (length > 0) memcpy(out->title, xcb_get_property_value(name_reply), length); + out->title[length] = 0; + } + free(name_reply); + } + return 1; +} + +// Returns 0 on success, 1 on connection error, 2 when no matching window exists. +static int aiscan_x11_resolve(const char *display, uint32_t requested_window, + uint32_t requested_pid, aiscan_x11_window *out, + uint32_t *screen_width, uint32_t *screen_height) { + int screen_number = 0; + xcb_connection_t *c = xcb_connect(display && display[0] ? display : NULL, &screen_number); + if (!c || xcb_connection_has_error(c)) { + if (c) xcb_disconnect(c); + return 1; + } + xcb_screen_t *screen = aiscan_screen(c, screen_number); + if (!screen) { + xcb_disconnect(c); + return 1; + } + *screen_width = screen->width_in_pixels; + *screen_height = screen->height_in_pixels; + xcb_atom_t pid_atom = aiscan_atom(c, "_NET_WM_PID"); + xcb_atom_t name_atom = aiscan_atom(c, "_NET_WM_NAME"); + + if (requested_window) { + int ok = aiscan_window_info(c, requested_window, pid_atom, name_atom, out); + xcb_disconnect(c); + return ok ? 0 : 2; + } + + xcb_atom_t list_atom = aiscan_atom(c, "_NET_CLIENT_LIST"); + xcb_get_property_reply_t *list = xcb_get_property_reply(c, + xcb_get_property(c, 0, screen->root, list_atom, XCB_ATOM_WINDOW, 0, UINT32_MAX), NULL); + if (!list) { + xcb_disconnect(c); + return 2; + } + int count = xcb_get_property_value_length(list) / (int)sizeof(xcb_window_t); + xcb_window_t *windows = (xcb_window_t *)xcb_get_property_value(list); + uint64_t best_area = 0; + aiscan_x11_window candidate; + memset(out, 0, sizeof(*out)); + for (int i = 0; i < count; i++) { + if (!aiscan_window_info(c, windows[i], pid_atom, name_atom, &candidate)) continue; + if (candidate.pid != requested_pid) continue; + uint64_t area = (uint64_t)candidate.width * candidate.height; + if (area > best_area) { + best_area = area; + *out = candidate; + } + } + free(list); + xcb_disconnect(c); + return out->window ? 0 : 2; +} +*/ +import "C" + +import ( + "context" + "fmt" + "math" + "os" + "strconv" + "strings" + "unsafe" +) + +func resolvePlatformTarget(_ context.Context, req captureRequest) (resolvedTarget, error) { + if req.WindowHandle > math.MaxUint32 { + return resolvedTarget{}, fmt.Errorf("X11 window ID 0x%x exceeds 32 bits", req.WindowHandle) + } + if strings.EqualFold(strings.TrimSpace(os.Getenv("XDG_SESSION_TYPE")), "wayland") { + return resolvedTarget{}, fmt.Errorf("Wayland capture is not supported; use an X11 session") + } + display := strings.TrimSpace(os.Getenv("DISPLAY")) + if display == "" { + return resolvedTarget{}, fmt.Errorf("DISPLAY is not set") + } + cDisplay := C.CString(display) + defer C.free(unsafe.Pointer(cDisplay)) + var info C.aiscan_x11_window + var screenWidth, screenHeight C.uint32_t + code := C.aiscan_x11_resolve( + cDisplay, + C.uint32_t(req.WindowHandle), + C.uint32_t(req.PID), + &info, + &screenWidth, + &screenHeight, + ) + if code == 1 { + return resolvedTarget{}, fmt.Errorf("connect to X11 display %s", display) + } + if req.Target == "desktop" { + return resolvedTarget{ + Info: TargetInfo{Kind: "desktop", Width: int(screenWidth), Height: int(screenHeight)}, + Native: nativeCaptureTarget{format: "x11grab", url: display}, + }, nil + } + if code != 0 { + if req.WindowHandle != 0 { + return resolvedTarget{}, fmt.Errorf("X11 window 0x%x is missing, hidden, or minimized", req.WindowHandle) + } + return resolvedTarget{}, fmt.Errorf("no visible X11 top-level window found for pid %d", req.PID) + } + handle := uint64(info.window) + return resolvedTarget{ + Info: TargetInfo{ + Kind: "window", + WindowHandle: fmt.Sprintf("0x%x", handle), + PID: int64(info.pid), + Title: C.GoString(&info.title[0]), + Width: int(info.width), + Height: int(info.height), + }, + Native: nativeCaptureTarget{ + format: "x11grab", + url: display, + options: map[string]string{ + "window_id": strconv.FormatUint(handle, 10), + }, + }, + }, nil +} diff --git a/tools/record/target_windows.go b/tools/record/target_windows.go new file mode 100644 index 00000000..c5e4d396 --- /dev/null +++ b/tools/record/target_windows.go @@ -0,0 +1,127 @@ +//go:build record_ffmpeg && cgo && windows + +package record + +import ( + "context" + "fmt" + "syscall" + "unsafe" + + "golang.org/x/sys/windows" +) + +var ( + user32 = windows.NewLazySystemDLL("user32.dll") + procEnumWindows = user32.NewProc("EnumWindows") + procGetWindowThreadProcessID = user32.NewProc("GetWindowThreadProcessId") + procIsWindow = user32.NewProc("IsWindow") + procIsWindowVisible = user32.NewProc("IsWindowVisible") + procIsIconic = user32.NewProc("IsIconic") + procGetWindowRect = user32.NewProc("GetWindowRect") + procGetWindowTextLengthW = user32.NewProc("GetWindowTextLengthW") + procGetWindowTextW = user32.NewProc("GetWindowTextW") + procGetSystemMetrics = user32.NewProc("GetSystemMetrics") + procGetForegroundWindow = user32.NewProc("GetForegroundWindow") +) + +type winRect struct { + Left int32 + Top int32 + Right int32 + Bottom int32 +} + +type winTargetInfo struct { + handle uint64 + pid int64 + title string + width int + height int +} + +func resolvePlatformTarget(_ context.Context, req captureRequest) (resolvedTarget, error) { + if req.Target == "desktop" { + width, _, _ := procGetSystemMetrics.Call(78) + height, _, _ := procGetSystemMetrics.Call(79) + return resolvedTarget{ + Info: TargetInfo{Kind: "desktop", Width: int(width), Height: int(height)}, + Native: nativeCaptureTarget{format: "gdigrab", url: "desktop"}, + }, nil + } + handle := req.WindowHandle + if handle == 0 { + resolved, err := findWindowByPID(uint32(req.PID)) + if err != nil { + return resolvedTarget{}, err + } + handle = resolved.handle + } + info, err := inspectWindow(uintptr(handle)) + if err != nil { + return resolvedTarget{}, err + } + return resolvedTarget{ + Info: TargetInfo{ + Kind: "window", + WindowHandle: fmt.Sprintf("0x%x", info.handle), + PID: info.pid, + Title: info.title, + Width: info.width, + Height: info.height, + }, + Native: nativeCaptureTarget{format: "gdigrab", url: fmt.Sprintf("hwnd=0x%x", info.handle)}, + }, nil +} + +func findWindowByPID(pid uint32) (winTargetInfo, error) { + var best winTargetInfo + callback := syscall.NewCallback(func(hwnd uintptr, _ uintptr) uintptr { + info, err := inspectWindow(hwnd) + if err != nil || uint32(info.pid) != pid { + return 1 + } + if info.width*info.height > best.width*best.height { + best = info + } + return 1 + }) + result, _, callErr := procEnumWindows.Call(callback, 0) + if result == 0 { + return winTargetInfo{}, fmt.Errorf("enumerate windows for pid %d: %w", pid, callErr) + } + if best.handle == 0 { + return winTargetInfo{}, fmt.Errorf("no visible non-minimized top-level window found for pid %d", pid) + } + return best, nil +} + +func inspectWindow(hwnd uintptr) (winTargetInfo, error) { + if ok, _, _ := procIsWindow.Call(hwnd); ok == 0 { + return winTargetInfo{}, fmt.Errorf("window handle 0x%x does not exist", hwnd) + } + if visible, _, _ := procIsWindowVisible.Call(hwnd); visible == 0 { + return winTargetInfo{}, fmt.Errorf("window 0x%x is not visible", hwnd) + } + if iconic, _, _ := procIsIconic.Call(hwnd); iconic != 0 { + return winTargetInfo{}, fmt.Errorf("window 0x%x is minimized", hwnd) + } + var rect winRect + if ok, _, callErr := procGetWindowRect.Call(hwnd, uintptr(unsafe.Pointer(&rect))); ok == 0 { + return winTargetInfo{}, fmt.Errorf("get window rect 0x%x: %w", hwnd, callErr) + } + width, height := int(rect.Right-rect.Left), int(rect.Bottom-rect.Top) + if width <= 0 || height <= 0 { + return winTargetInfo{}, fmt.Errorf("window 0x%x has invalid dimensions %dx%d", hwnd, width, height) + } + var pid uint32 + procGetWindowThreadProcessID.Call(hwnd, uintptr(unsafe.Pointer(&pid))) + length, _, _ := procGetWindowTextLengthW.Call(hwnd) + var title string + if length > 0 { + buf := make([]uint16, length+1) + procGetWindowTextW.Call(hwnd, uintptr(unsafe.Pointer(&buf[0])), length+1) + title = windows.UTF16ToString(buf) + } + return winTargetInfo{handle: uint64(hwnd), pid: int64(pid), title: title, width: width, height: height}, nil +} diff --git a/tools/record/target_windows_integration_test.go b/tools/record/target_windows_integration_test.go new file mode 100644 index 00000000..759cef09 --- /dev/null +++ b/tools/record/target_windows_integration_test.go @@ -0,0 +1,46 @@ +//go:build record_ffmpeg && record_integration && cgo && windows + +package record + +import ( + "context" + "path/filepath" + "testing" + "time" +) + +func TestNativeWindowHandleAndPIDCapture(t *testing.T) { + hwnd, _, _ := procGetForegroundWindow.Call() + if hwnd == 0 { + t.Log("no foreground window; desktop integration still covered") + return + } + backend := newPlatformBackend() + byHandle, err := backend.Resolve(context.Background(), captureRequest{Target: "window", WindowHandle: uint64(hwnd), FPS: 8}) + if err != nil { + t.Logf("foreground window is not stable enough to capture: %v", err) + return + } + if byHandle.Info.PID <= 0 { + t.Fatal("resolved window has no PID") + } + if _, err := backend.Screenshot(context.Background(), byHandle); err != nil { + t.Fatalf("screenshot by handle: %v", err) + } + byPID, err := backend.Resolve(context.Background(), captureRequest{Target: "window", PID: byHandle.Info.PID, FPS: 8}) + if err != nil { + // The foreground window belongs to another interactive process and may + // close, minimize, or replace its top-level HWND between the two calls. + t.Logf("foreground process is no longer capturable by PID: %v", err) + return + } + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + media, err := backend.Record(ctx, byPID, filepath.Join(t.TempDir(), "window.mp4"), 8) + if err != nil { + t.Fatalf("record by PID: %v", err) + } + if media.Frames == 0 { + t.Fatal("window recording produced no frames") + } +} diff --git a/tools/record/tool.go b/tools/record/tool.go new file mode 100644 index 00000000..627f06ec --- /dev/null +++ b/tools/record/tool.go @@ -0,0 +1,157 @@ +package record + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "sync" + + aop "github.com/chainreactors/aiscan/aop" + "github.com/chainreactors/aiscan/core/tool" + "github.com/chainreactors/aiscan/pkg/imageutil" +) + +type Tool struct { + workDir string + outputDir string + maxConcurrent int + backend captureBackend + + mu sync.RWMutex + sessions map[string]*recordingSession + closed bool +} + +func New(workDir, outputDir string, maxConcurrent int, backend captureBackend) *Tool { + if maxConcurrent <= 0 { + maxConcurrent = defaultMaxConcurrent + } + if maxConcurrent > maxConcurrentLimit { + maxConcurrent = maxConcurrentLimit + } + return &Tool{ + workDir: workDir, + outputDir: outputDir, + maxConcurrent: maxConcurrent, + backend: backend, + sessions: make(map[string]*recordingSession), + } +} + +func (t *Tool) Name() string { return "record" } + +func (t *Tool) Description() string { + return "Capture desktop or application-window screenshots and H.264 MP4 recordings. Supports synchronous duration recording and asynchronous start/stop/status sessions." +} + +func (t *Tool) Definition() *tool.Definition { + return tool.Def(t.Name(), t.Description(), Args{}) +} + +func (t *Tool) Execute(ctx context.Context, arguments string) (*tool.Result, error) { + args, err := tool.ParseArgs[Args](arguments) + if err != nil { + return nil, err + } + action := strings.ToLower(strings.TrimSpace(args.Action)) + if action == "" { + return nil, fmt.Errorf("action is required") + } + if action != "status" && t.isClosed() { + return nil, fmt.Errorf("record tool is closed") + } + + switch action { + case "screenshot": + return t.screenshot(ctx, args) + case "record": + duration, err := normalizeDuration(args.DurationSeconds, true) + if err != nil { + return nil, err + } + session, err := t.start(ctx, args, duration) + if err != nil { + return nil, err + } + return t.waitResult(ctx, session) + case "start": + duration, err := normalizeDuration(args.DurationSeconds, false) + if err != nil { + return nil, err + } + session, err := t.start(ctx, args, duration) + if err != nil { + return nil, err + } + return jsonResult(session.snapshot()) + case "stop": + return t.stop(ctx, strings.TrimSpace(args.RecordingID)) + case "status": + return t.status(strings.TrimSpace(args.RecordingID)) + default: + return nil, fmt.Errorf("unsupported action %q", args.Action) + } +} + +func (t *Tool) screenshot(ctx context.Context, args Args) (*tool.Result, error) { + if t.backend == nil { + return nil, fmt.Errorf("capture backend is unavailable") + } + req, err := normalizeCaptureArgs(args) + if err != nil { + return nil, err + } + target, err := callBackendResolve(t.backend, ctx, req) + if err != nil { + return nil, fmt.Errorf("resolve capture target: %w", err) + } + img, err := callBackendScreenshot(t.backend, ctx, target) + if err != nil { + return nil, fmt.Errorf("capture screenshot: %w", err) + } + if img == nil { + return nil, fmt.Errorf("capture screenshot: backend returned no image") + } + if err := ctx.Err(); err != nil { + return nil, err + } + + target.Info.Width = img.Bounds().Dx() + target.Info.Height = img.Bounds().Dy() + path, err := t.outputPath(ctx, args.Output, "screenshot-"+newID(), ".png") + if err != nil { + return nil, err + } + data := imageutil.EncodePNG(img) + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + return nil, fmt.Errorf("create screenshot directory: %w", err) + } + if err := os.WriteFile(path, data, 0o644); err != nil { + return nil, fmt.Errorf("write screenshot: %w", err) + } + preview, err := imageutil.OptimizeImage(img) + if err != nil { + return nil, fmt.Errorf("prepare screenshot preview: %w", err) + } + meta := struct { + Action string `json:"action"` + Target TargetInfo `json:"target"` + Output string `json:"output"` + Bytes int `json:"bytes"` + MimeType string `json:"mime_type"` + }{"screenshot", target.Info, path, len(data), "image/png"} + text, _ := json.MarshalIndent(meta, "", " ") + return &tool.Result{Output: []*aop.Content{ + aop.Text(string(text)), + aop.MediaData("image", preview.MimeType, filepath.Base(path), preview.Data), + }}, nil +} + +func (t *Tool) isClosed() bool { + t.mu.RLock() + defer t.mu.RUnlock() + return t.closed +} diff --git a/tools/record/tool_test.go b/tools/record/tool_test.go new file mode 100644 index 00000000..705af4fa --- /dev/null +++ b/tools/record/tool_test.go @@ -0,0 +1,397 @@ +package record + +import ( + "context" + "encoding/json" + "errors" + "image" + "image/color" + "math" + "os" + "path/filepath" + "runtime" + "strings" + "sync" + "testing" + "time" + + coretool "github.com/chainreactors/aiscan/core/tool" +) + +type fakeBackend struct { + mu sync.Mutex + startOnce sync.Once + started chan struct{} + resolveErr error + screenshot image.Image + screenshotErr error + nilScreenshot bool + record func(context.Context, string) (mediaInfo, error) +} + +func (b *fakeBackend) Resolve(_ context.Context, req captureRequest) (resolvedTarget, error) { + if b.resolveErr != nil { + return resolvedTarget{}, b.resolveErr + } + info := TargetInfo{Kind: req.Target, PID: req.PID, Width: 320, Height: 240} + if req.WindowHandle != 0 { + info.WindowHandle = "0x1" + } + return resolvedTarget{Info: info}, nil +} + +func (b *fakeBackend) Screenshot(context.Context, resolvedTarget) (image.Image, error) { + if b.screenshotErr != nil { + return nil, b.screenshotErr + } + if b.nilScreenshot { + return nil, nil + } + if b.screenshot != nil { + return b.screenshot, nil + } + img := image.NewRGBA(image.Rect(0, 0, 320, 240)) + img.Set(1, 1, color.RGBA{R: 255, A: 255}) + return img, nil +} + +func (b *fakeBackend) Record(ctx context.Context, _ resolvedTarget, output string, _ int) (mediaInfo, error) { + b.mu.Lock() + if b.started != nil { + b.startOnce.Do(func() { close(b.started) }) + } + b.mu.Unlock() + if b.record != nil { + return b.record(ctx, output) + } + <-ctx.Done() + if err := os.WriteFile(output, []byte("fake-mp4"), 0o644); err != nil { + return mediaInfo{}, err + } + return mediaInfo{Width: 320, Height: 240, Frames: 3}, nil +} + +func TestScreenshotReturnsImageAndPath(t *testing.T) { + dir := t.TempDir() + tool := New(dir, filepath.Join(dir, "record"), 4, &fakeBackend{}) + result, err := tool.Execute(context.Background(), `{"action":"screenshot"}`) + if err != nil { + t.Fatal(err) + } + if !coretool.ResultHasImages(result) { + t.Fatal("screenshot result should contain an image") + } + text := coretool.ResultText(result) + var meta struct { + Output string `json:"output"` + } + if err := json.Unmarshal([]byte(text), &meta); err != nil { + t.Fatal(err) + } + if _, err := os.Stat(meta.Output); err != nil { + t.Fatalf("screenshot output: %v", err) + } +} + +func TestAsyncStartStopStatus(t *testing.T) { + dir := t.TempDir() + backend := &fakeBackend{started: make(chan struct{})} + tool := New(dir, filepath.Join(dir, "record"), 4, backend) + result, err := tool.Execute(context.Background(), `{"action":"start","target":"window","window_handle":"0x1"}`) + if err != nil { + t.Fatal(err) + } + var info SessionInfo + if err := json.Unmarshal([]byte(coretool.ResultText(result)), &info); err != nil { + t.Fatal(err) + } + select { + case <-backend.started: + case <-time.After(time.Second): + t.Fatal("recording did not start") + } + stopArgs := `{"action":"stop","recording_id":"` + info.RecordingID + `"}` + result, err = tool.Execute(context.Background(), stopArgs) + if err != nil { + t.Fatal(err) + } + if err := json.Unmarshal([]byte(coretool.ResultText(result)), &info); err != nil { + t.Fatal(err) + } + if info.State != "completed" || info.Frames != 3 || info.Bytes == 0 { + t.Fatalf("unexpected final info: %+v", info) + } + statusArgs := `{"action":"status","recording_id":"` + info.RecordingID + `"}` + result, err = tool.Execute(context.Background(), statusArgs) + if err != nil || !strings.Contains(coretool.ResultText(result), `"state": "completed"`) { + t.Fatalf("status result=%q err=%v", coretool.ResultText(result), err) + } +} + +func TestConcurrencyLimit(t *testing.T) { + dir := t.TempDir() + tool := New(dir, filepath.Join(dir, "record"), 2, &fakeBackend{}) + for i := 0; i < 2; i++ { + if _, err := tool.Execute(context.Background(), `{"action":"start"}`); err != nil { + t.Fatal(err) + } + } + if _, err := tool.Execute(context.Background(), `{"action":"start"}`); err == nil || !strings.Contains(err.Error(), "concurrency limit 2") { + t.Fatalf("expected concurrency error, got %v", err) + } + tool.Close() +} + +func TestValidation(t *testing.T) { + tool := New(t.TempDir(), t.TempDir(), 4, &fakeBackend{}) + tests := []string{ + `{"action":"record"}`, + `{"action":"screenshot","target":"window"}`, + `{"action":"screenshot","target":"window","pid":1,"window_handle":"1"}`, + `{"action":"screenshot","target":"window","pid":4294967296}`, + `{"action":"screenshot","fps":61}`, + `{"action":"stop"}`, + } + for _, input := range tests { + if _, err := tool.Execute(context.Background(), input); err == nil { + t.Errorf("expected error for %s", input) + } + } +} + +func TestRelativeOutputUsesInvocationWorkDir(t *testing.T) { + dir := t.TempDir() + tool := New("ignored", filepath.Join(dir, "default"), 4, &fakeBackend{}) + ctx := coretool.ContextWithInvocation(context.Background(), coretool.Invocation{WorkDir: dir}) + result, err := tool.Execute(ctx, `{"action":"screenshot","output":"shots/test.png"}`) + if err != nil { + t.Fatal(err) + } + var meta struct { + Output string `json:"output"` + } + if err := json.Unmarshal([]byte(coretool.ResultText(result)), &meta); err != nil { + t.Fatal(err) + } + if meta.Output != filepath.Join(dir, "shots", "test.png") { + t.Fatalf("result path = %s", meta.Output) + } +} + +func TestDefaultOutputUsesInvocationRecordDir(t *testing.T) { + dir := t.TempDir() + tool := New("ignored", filepath.Join(t.TempDir(), "fallback"), 4, &fakeBackend{}) + ctx := coretool.ContextWithInvocation(context.Background(), coretool.Invocation{WorkDir: dir}) + result, err := tool.Execute(ctx, `{"action":"screenshot"}`) + if err != nil { + t.Fatal(err) + } + var meta struct { + Output string `json:"output"` + } + if err := json.Unmarshal([]byte(coretool.ResultText(result)), &meta); err != nil { + t.Fatal(err) + } + wantDir := filepath.Join(dir, ".aiscan", "record") + if filepath.Dir(meta.Output) != wantDir { + t.Fatalf("result directory = %s, want %s", filepath.Dir(meta.Output), wantDir) + } +} + +func TestSynchronousRecordCompletesAfterDuration(t *testing.T) { + dir := t.TempDir() + recorder := New(dir, filepath.Join(dir, "record"), 1, &fakeBackend{}) + result, err := recorder.Execute(context.Background(), `{"action":"record","duration_seconds":0.01}`) + if err != nil { + t.Fatal(err) + } + var info SessionInfo + if err := json.Unmarshal([]byte(coretool.ResultText(result)), &info); err != nil { + t.Fatal(err) + } + if info.State != sessionCompleted || info.Frames != 3 || info.Bytes == 0 { + t.Fatalf("unexpected final info: %+v", info) + } + if info.DurationMS < 1 { + t.Fatalf("duration_ms = %d, want positive duration", info.DurationMS) + } + if !coretool.ResultHasMedia(result) { + t.Fatal("record result should contain video media") + } + media := result.Output[1].GetMedia() + if media == nil || media.Kind != "video" || media.Resource.GetMediaType() != "video/mp4" || media.Resource.GetUri() == "" { + t.Fatalf("video media = %+v", media) + } + if filepath.IsAbs(media.Resource.GetUri()) { + t.Fatalf("video URI should be relative to the invocation workdir: %q", media.Resource.GetUri()) + } +} + +func TestBackendFailureReturnsStructuredToolError(t *testing.T) { + backend := &fakeBackend{record: func(context.Context, string) (mediaInfo, error) { + return mediaInfo{}, errors.New("encoder failed") + }} + recorder := New(t.TempDir(), t.TempDir(), 1, backend) + result, err := recorder.Execute(context.Background(), `{"action":"record","duration_seconds":1}`) + if err != nil { + t.Fatal(err) + } + if !result.GetIsError() { + t.Fatalf("result should be marked as an error: %s", coretool.ResultText(result)) + } + var info SessionInfo + if err := json.Unmarshal([]byte(coretool.ResultText(result)), &info); err != nil { + t.Fatal(err) + } + if info.State != sessionFailed || !strings.Contains(info.Error, "encoder failed") { + t.Fatalf("unexpected failed session: %+v", info) + } +} + +func TestBackendPanicIsContainedInSession(t *testing.T) { + backend := &fakeBackend{record: func(context.Context, string) (mediaInfo, error) { + panic("native crash") + }} + recorder := New(t.TempDir(), t.TempDir(), 1, backend) + result, err := recorder.Execute(context.Background(), `{"action":"record","duration_seconds":1}`) + if err != nil { + t.Fatal(err) + } + if !result.GetIsError() || !strings.Contains(coretool.ResultText(result), "capture backend panicked: native crash") { + t.Fatalf("panic was not converted to a session error: %s", coretool.ResultText(result)) + } +} + +func TestConcurrentStartAdmissionIsAtomic(t *testing.T) { + recorder := New(t.TempDir(), t.TempDir(), 1, &fakeBackend{}) + t.Cleanup(recorder.Close) + const attempts = 16 + var wg sync.WaitGroup + errs := make(chan error, attempts) + for range attempts { + wg.Add(1) + go func() { + defer wg.Done() + _, err := recorder.Execute(context.Background(), `{"action":"start"}`) + errs <- err + }() + } + wg.Wait() + close(errs) + successes := 0 + for err := range errs { + if err == nil { + successes++ + continue + } + if !strings.Contains(err.Error(), "concurrency limit 1") { + t.Errorf("unexpected start error: %v", err) + } + } + if successes != 1 { + t.Fatalf("successful starts = %d, want 1", successes) + } +} + +func TestDuplicateActiveOutputIsRejected(t *testing.T) { + dir := t.TempDir() + recorder := New(dir, filepath.Join(dir, "record"), 2, &fakeBackend{}) + t.Cleanup(recorder.Close) + if _, err := recorder.Execute(context.Background(), `{"action":"start","output":"same.mp4"}`); err != nil { + t.Fatal(err) + } + if _, err := recorder.Execute(context.Background(), `{"action":"start","output":"same.mp4"}`); err == nil || !strings.Contains(err.Error(), "already used") { + t.Fatalf("expected duplicate output error, got %v", err) + } +} + +func TestSessionHistoryIsBounded(t *testing.T) { + backend := &fakeBackend{record: func(_ context.Context, output string) (mediaInfo, error) { + if err := os.WriteFile(output, []byte("mp4"), 0o644); err != nil { + return mediaInfo{}, err + } + return mediaInfo{Frames: 1, Width: 2, Height: 2}, nil + }} + recorder := New(t.TempDir(), t.TempDir(), 1, backend) + for range maxSessionHistory + 5 { + session, err := recorder.start(context.Background(), Args{}, 0) + if err != nil { + t.Fatal(err) + } + <-session.done + } + recorder.mu.RLock() + count := len(recorder.sessions) + recorder.mu.RUnlock() + if count != maxSessionHistory { + t.Fatalf("retained sessions = %d, want %d", count, maxSessionHistory) + } +} + +func TestClosedToolRejectsNewCaptureButKeepsStatus(t *testing.T) { + recorder := New(t.TempDir(), t.TempDir(), 1, &fakeBackend{}) + recorder.Close() + if _, err := recorder.Execute(context.Background(), `{"action":"screenshot"}`); err == nil || !strings.Contains(err.Error(), "closed") { + t.Fatalf("expected closed error, got %v", err) + } + if _, err := recorder.Execute(context.Background(), `{"action":"status"}`); err != nil { + t.Fatalf("status after close: %v", err) + } +} + +func TestUnavailableOrInvalidScreenshotBackend(t *testing.T) { + for name, backend := range map[string]captureBackend{ + "nil backend": nil, + "backend error": &fakeBackend{screenshotErr: errNilImage}, + "nil image": &fakeBackend{nilScreenshot: true}, + } { + t.Run(name, func(t *testing.T) { + recorder := New(t.TempDir(), t.TempDir(), 1, backend) + _, err := recorder.Execute(context.Background(), `{"action":"screenshot"}`) + if err == nil { + t.Fatal("expected screenshot error") + } + }) + } +} + +var errNilImage = errors.New("no frame") + +func TestNormalizeDuration(t *testing.T) { + tests := []struct { + name string + seconds float64 + required bool + wantErr string + }{ + {name: "optional zero", seconds: 0}, + {name: "required zero", seconds: 0, required: true, wantErr: "greater than zero"}, + {name: "negative", seconds: -1, wantErr: "cannot be negative"}, + {name: "nan", seconds: math.NaN(), wantErr: "finite"}, + {name: "infinity", seconds: math.Inf(1), wantErr: "finite"}, + {name: "overflow", seconds: float64(math.MaxInt64), wantErr: "too large"}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := normalizeDuration(tt.seconds, tt.required) + if tt.wantErr == "" && err != nil { + t.Fatal(err) + } + if tt.wantErr != "" && (err == nil || !strings.Contains(err.Error(), tt.wantErr)) { + t.Fatalf("error = %v, want substring %q", err, tt.wantErr) + } + }) + } +} + +func TestSamePathUsesPlatformCaseSemantics(t *testing.T) { + a := filepath.Join(t.TempDir(), "Capture.mp4") + b := filepath.Join(filepath.Dir(a), "capture.mp4") + got := samePath(a, b) + if runtime.GOOS == "windows" && !got { + t.Fatal("Windows paths should compare case-insensitively") + } + if runtime.GOOS != "windows" && got { + t.Fatal("non-Windows paths should compare case-sensitively") + } +} diff --git a/tools/record/types.go b/tools/record/types.go new file mode 100644 index 00000000..228c3149 --- /dev/null +++ b/tools/record/types.go @@ -0,0 +1,77 @@ +package record + +import ( + "context" + "image" + "time" +) + +const ( + defaultFPS = 30 + maxSessionHistory = 128 + sessionStopTimeout = 10 * time.Second + + sessionStarting = "starting" + sessionRecording = "recording" + sessionStopping = "stopping" + sessionCompleted = "completed" + sessionFailed = "failed" +) + +type Args struct { + Action string `json:"action" jsonschema:"description=Operation to perform.,enum=screenshot,enum=record,enum=start,enum=stop,enum=status"` + Target string `json:"target,omitempty" jsonschema:"description=Capture target. Defaults to desktop.,enum=desktop,enum=window"` + WindowHandle string `json:"window_handle,omitempty" jsonschema:"description=Windows HWND or X11 window ID in decimal or 0x hexadecimal form"` + PID int64 `json:"pid,omitempty" jsonschema:"description=Process ID used to resolve the largest visible top-level window"` + DurationSeconds float64 `json:"duration_seconds,omitempty" jsonschema:"description=Duration for action=record; optional auto-stop duration for action=start"` + FPS int `json:"fps,omitempty" jsonschema:"description=Capture frame rate 1-60 (default 30),minimum=1,maximum=60"` + Output string `json:"output,omitempty" jsonschema:"description=Output path. Relative paths resolve against the invocation working directory"` + RecordingID string `json:"recording_id,omitempty" jsonschema:"description=Recording identifier required by stop and optional for status"` +} + +type TargetInfo struct { + Kind string `json:"kind"` + WindowHandle string `json:"window_handle,omitempty"` + PID int64 `json:"pid,omitempty"` + Title string `json:"title,omitempty"` + Width int `json:"width,omitempty"` + Height int `json:"height,omitempty"` +} + +type SessionInfo struct { + RecordingID string `json:"recording_id"` + State string `json:"state"` + Target TargetInfo `json:"target"` + Output string `json:"output"` + FPS int `json:"fps"` + Frames int64 `json:"frames,omitempty"` + Bytes int64 `json:"bytes,omitempty"` + StartedAt *time.Time `json:"started_at,omitempty"` + EndedAt *time.Time `json:"ended_at,omitempty"` + DurationMS int64 `json:"duration_ms,omitempty"` + Error string `json:"error,omitempty"` +} + +type captureRequest struct { + Target string + WindowHandle uint64 + PID int64 + FPS int +} + +type resolvedTarget struct { + Info TargetInfo + Native any +} + +type mediaInfo struct { + Width int + Height int + Frames int64 +} + +type captureBackend interface { + Resolve(context.Context, captureRequest) (resolvedTarget, error) + Screenshot(context.Context, resolvedTarget) (image.Image, error) + Record(context.Context, resolvedTarget, string, int) (mediaInfo, error) +} diff --git a/tools/register_command_test.go b/tools/register_command_test.go index 3fc4fd3b..0d2a2844 100644 --- a/tools/register_command_test.go +++ b/tools/register_command_test.go @@ -571,7 +571,7 @@ http: Name: "neutron/custom-poc-filter-json", Tool: "neutron", Args: []string{"-i", httpServer.URL, "-t", templateFile, "--tags", "regression", "-s", "high", "-j"}, Check: func(t *testing.T, result functionalResult) { - requireOutputContains(t, result, `"matched":true`, `"template":"regression-marker"`) + requireOutputContains(t, result, `"matched":true`, `"template_id":"regression-marker"`) requireEvent(t, result, "neutron", toolpb.ArtifactKindVuln, nil) }, }, diff --git a/tools/scan/adapter.go b/tools/scan/adapter.go index c5a56a6a..812a6a20 100644 --- a/tools/scan/adapter.go +++ b/tools/scan/adapter.go @@ -6,7 +6,10 @@ import ( "net/url" "strings" + toolpb "github.com/chainreactors/aiscan/aop/tool" + "github.com/chainreactors/aiscan/core/output" "github.com/chainreactors/aiscan/tools/scan/engine" + "github.com/chainreactors/aiscan/tools/toolargs" sdktypes "github.com/chainreactors/sdk/pkg/types" sdkzombie "github.com/chainreactors/sdk/zombie" "github.com/chainreactors/utils" @@ -170,7 +173,16 @@ func (c *Command) runPOCCapability(ctx context.Context, flags flags, input targe if result == nil || !result.Matched() { continue } - emit(lootEvent(capNeutronPOC, vulnLoot(result.TemplateResult(target.Target)))) + record := result.TemplateResult(target.Target) + resultID := toolargs.ArtifactResultID("neutron", toolpb.ArtifactKindVuln, target.Target, record) + loot := bindLoot(vulnLoot(record), resultID, "neutron") + emit(artifactLootEvent(capNeutronPOC, loot, output.ArtifactResult{ + ResultID: resultID, + Tool: "neutron", + Kind: toolpb.ArtifactKindVuln, + Target: target.Target, + Data: record, + })) } } @@ -189,7 +201,12 @@ func deriveServiceResult(profile profile, source string, result *parsers.GOGORes emit(targetEvent(source, "", newWebTarget("", target, ""))) } if len(fingers) > 0 { - emit(lootEvent(source, fingerprintLoot(target, parsers.NormalizeNames(fingers), result.Frameworks.IsFocus()))) + resultID := toolargs.ArtifactResultID("gogo", toolpb.ArtifactKindService, result.GetTarget(), result) + emit(lootEvent(source, bindLoot( + fingerprintLoot(target, parsers.NormalizeNames(fingers), result.Frameworks.IsFocus()), + resultID, + "gogo", + ))) } if len(fingers) > 0 || profile.AllowBroadPOC { emit(targetEvent(source, "", newPOCTarget("", target, fingers))) @@ -217,7 +234,12 @@ func deriveWebProbeResult(profile profile, source string, result *parsers.SprayR } fingers := parsers.FrameworkNames(result.Frameworks) if len(fingers) > 0 { - emit(lootEvent(source, fingerprintLoot(result.UrlString, parsers.NormalizeNames(fingers), result.Frameworks.IsFocus()))) + resultID := toolargs.ArtifactResultID("spray", toolpb.ArtifactKindWeb, result.UrlString, result) + emit(lootEvent(source, bindLoot( + fingerprintLoot(result.UrlString, parsers.NormalizeNames(fingers), result.Frameworks.IsFocus()), + resultID, + "spray", + ))) } if result.Status > 0 && (len(fingers) > 0 || profile.AllowBroadPOC) { emit(targetEvent(source, "", newPOCTarget("", result.UrlString, fingers))) @@ -305,7 +327,16 @@ func deriveWeakpassResult(source string, result *parsers.ZombieResult, emit func if result == nil { return } - emit(lootEvent(source, weakpassLoot(result))) + target := result.Address() + resultID := toolargs.ArtifactResultID("zombie", toolpb.ArtifactKindWeakpass, target, result) + loot := bindLoot(weakpassLoot(result), resultID, "zombie") + emit(artifactLootEvent(source, loot, output.ArtifactResult{ + ResultID: resultID, + Tool: "zombie", + Kind: toolpb.ArtifactKindWeakpass, + Target: target, + Data: result, + })) } func zombieTargetFromGogo(result *parsers.GOGOResult) (sdkzombie.Target, bool) { diff --git a/tools/scan/capability_katana.go b/tools/scan/capability_katana.go index 74699eb0..ce5ba249 100644 --- a/tools/scan/capability_katana.go +++ b/tools/scan/capability_katana.go @@ -9,8 +9,11 @@ import ( "strings" "sync" + browserutil "github.com/chainreactors/aiscan/pkg/browser" "github.com/projectdiscovery/gologger" "github.com/projectdiscovery/gologger/levels" + "github.com/projectdiscovery/katana/pkg/engine" + "github.com/projectdiscovery/katana/pkg/engine/headless" "github.com/projectdiscovery/katana/pkg/engine/standard" katanaoutput "github.com/projectdiscovery/katana/pkg/output" katanatypes "github.com/projectdiscovery/katana/pkg/types" @@ -85,6 +88,29 @@ func runKatanaCrawl(ctx context.Context, c *Command, e event, depth int, jsMode var mu sync.Mutex seen := make(map[string]struct{}) + handleResult := func(r *katanaoutput.Result) { + if r == nil || r.Request == nil || r.Request.URL == "" { + return + } + discoveredURL := r.Request.URL + + if seedRDN != "" && !sameRootDomain(discoveredURL, seedRDN) { + return + } + if strings.TrimRight(strings.ToLower(discoveredURL), "/") == seedNorm { + return + } + + mu.Lock() + if _, dup := seen[discoveredURL]; dup { + mu.Unlock() + return + } + seen[discoveredURL] = struct{}{} + mu.Unlock() + + emit(targetEvent(source, wt.Raw, newWebTarget(wt.Raw, discoveredURL, wt.HostHeader))) + } options := &katanatypes.Options{ MaxDepth: depth, @@ -95,36 +121,32 @@ func runKatanaCrawl(ctx context.Context, c *Command, e event, depth int, jsMode Silent: true, ScrapeJSResponses: jsMode, ScrapeJSLuiceResponses: jsMode, + Headless: jsMode, Timeout: 10, + TimeStable: 1, + MaxFailureCount: 10, + PageLoadStrategy: "heuristic", + DOMWaitTime: 5, Concurrency: 10, Parallelism: 10, OnResult: func(r katanaoutput.Result) { - if r.Request == nil || r.Request.URL == "" { - return - } - discoveredURL := r.Request.URL - - if seedRDN != "" && !sameRootDomain(discoveredURL, seedRDN) { - return - } - if strings.TrimRight(strings.ToLower(discoveredURL), "/") == seedNorm { - return - } - - mu.Lock() - if _, dup := seen[discoveredURL]; dup { - mu.Unlock() - return - } - seen[discoveredURL] = struct{}{} - mu.Unlock() - - emit(targetEvent(source, wt.Raw, newWebTarget(wt.Raw, discoveredURL, wt.HostHeader))) + handleResult(&r) }, } if c.Proxy != "" { options.Proxy = c.Proxy } + if jsMode { + binary, err := browserutil.Discover() + if err != nil { + emitError(emit, source, "katana browser discovery: %v", err) + return + } + if binary.Path != "" { + options.SystemChromePath = binary.Path + options.UseInstalledChrome = true + } + } gologger.DefaultLogger.SetMaxLevel(levels.LevelSilent) crawlerOptions, err := katanatypes.NewCrawlerOptions(options) @@ -133,13 +155,18 @@ func runKatanaCrawl(ctx context.Context, c *Command, e event, depth int, jsMode emitError(emit, source, "katana init %s: %v", wt.URL, err) return } - crawlerOptions.OutputWriter = &silentWriter{} + crawlerOptions.OutputWriter = &scanResultWriter{onResult: handleResult} defer func() { crawlerOptions.Close() gologger.DefaultLogger.SetMaxLevel(levels.LevelWarning) }() - crawler, err := standard.New(crawlerOptions) + var crawler engine.Engine + if jsMode { + crawler, err = headless.New(crawlerOptions) + } else { + crawler, err = standard.New(crawlerOptions) + } if err != nil { emitError(emit, source, "katana create %s: %v", wt.URL, err) return @@ -175,8 +202,15 @@ func sameRootDomain(rawURL, rdn string) bool { return host == rdn || strings.HasSuffix(host, "."+rdn) } -type silentWriter struct{} +type scanResultWriter struct { + onResult func(*katanaoutput.Result) +} -func (w *silentWriter) Close() error { return nil } -func (w *silentWriter) Write(_ *katanaoutput.Result) error { return nil } -func (w *silentWriter) WriteErr(_ *katanaoutput.Error) error { return nil } +func (w *scanResultWriter) Close() error { return nil } +func (w *scanResultWriter) Write(result *katanaoutput.Result) error { + if w.onResult != nil { + w.onResult(result) + } + return nil +} +func (w *scanResultWriter) WriteErr(_ *katanaoutput.Error) error { return nil } diff --git a/tools/scan/capability_katana_test.go b/tools/scan/capability_katana_test.go index b26dbd89..a17fb160 100644 --- a/tools/scan/capability_katana_test.go +++ b/tools/scan/capability_katana_test.go @@ -4,8 +4,15 @@ package scan import ( "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" "testing" "time" + + browserutil "github.com/chainreactors/aiscan/pkg/browser" ) func TestKatanaProfileExtender(t *testing.T) { @@ -58,3 +65,92 @@ func TestRunKatanaCrawlEmitsTargets(t *testing.T) { } t.Logf("katana discovered %d web targets from example.com (depth=1)", targets) } + +func TestE2EKatanaDeepRendersAuthenticatedSPA(t *testing.T) { + binary, err := browserutil.Discover() + if err != nil { + t.Fatalf("discover browser: %v", err) + } + if binary.Path == "" { + t.Skip("no system browser available; CI installs Chrome for this test") + } + t.Setenv(browserutil.PathEnv, binary.Path) + + const ( + sessionToken = "scan-session-73" + workspacePath = "/workspace/session-73?view=assets" + ) + var authenticatedWorkspaceHits atomic.Int32 + + mux := http.NewServeMux() + mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/" { + http.NotFound(w, r) + return + } + w.Header().Set("Content-Type", "text/html; charset=utf-8") + fmt.Fprint(w, `
Loading assets...
`) + }) + mux.HandleFunc("/api/session", func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost || r.Header.Get("X-CSRF-Token") != "scan-e2e" { + http.Error(w, "invalid session request", http.StatusForbidden) + return + } + w.Header().Set("Content-Type", "application/json") + fmt.Fprintf(w, `{"token":%q,"next":%q}`, sessionToken, workspacePath) + }) + mux.HandleFunc("/workspace/session-73", func(w http.ResponseWriter, r *http.Request) { + cookie, err := r.Cookie("scan_session") + if err != nil || cookie.Value != sessionToken { + http.Error(w, "authentication required", http.StatusUnauthorized) + return + } + authenticatedWorkspaceHits.Add(1) + fmt.Fprint(w, `Asset detail`) + }) + mux.HandleFunc("/assets/detail", func(w http.ResponseWriter, _ *http.Request) { + fmt.Fprint(w, "asset detail") + }) + srv := httptest.NewServer(mux) + defer srv.Close() + + cmd := &Command{} + e := targetEvent(capSprayCheck, srv.URL, newWebTarget(srv.URL, srv.URL, "")) + var emitted []event + ctx, cancel := context.WithTimeout(context.Background(), 75*time.Second) + defer cancel() + runKatanaCrawl(ctx, cmd, e, 2, true, func(ev event) { + emitted = append(emitted, ev) + }) + + for _, ev := range emitted { + if ev.Kind == eventError { + t.Fatalf("katana_deep emitted error: %s", ev.Error.Message) + } + } + if authenticatedWorkspaceHits.Load() == 0 { + t.Fatal("katana_deep browser never reached the authenticated workspace") + } + for _, ev := range emitted { + if ev.Kind != eventTarget { + continue + } + wt, ok := ev.Target.(webTarget) + if ok && strings.Contains(wt.URL, "/workspace/session-73") { + return + } + } + t.Fatal("katana_deep did not emit the browser-only workspace route") +} diff --git a/tools/scan/collector.go b/tools/scan/collector.go index fc4fd7b5..2b9729f0 100644 --- a/tools/scan/collector.go +++ b/tools/scan/collector.go @@ -26,6 +26,7 @@ type collector struct { stats *statsCollector gogoResults []*parsers.GOGOResult sprayResults []sprayObservation + artifacts []output.ArtifactResult loots []output.Loot errors []string trace []string @@ -128,6 +129,9 @@ func (c *collector) recordLootEvent(event event) { return } loot := *event.Loot + if event.Artifact != nil { + c.artifacts = append(c.artifacts, *event.Artifact) + } switch loot.Kind { case output.LootFingerprint: fingers := loot.Tags diff --git a/tools/scan/command.go b/tools/scan/command.go index d398eb4f..c0ceeff1 100644 --- a/tools/scan/command.go +++ b/tools/scan/command.go @@ -113,7 +113,7 @@ func (c *Command) execute(ctx context.Context, args []string, stream io.Writer) defer restoreDebug() c.Logger.Debugf("scan debug enabled") } - profile, err := profileForMode(flags.Mode) + profile, err := profileForFlags(flags) if err != nil { return "", nil, fmt.Errorf("scan: %w", err) } @@ -193,27 +193,47 @@ func (c *Command) emitStructuredData(ctx context.Context, result *output.ScanRes } for _, service := range result.GOGO { if service != nil { - c.EmitArtifactCtx(ctx, "gogo", toolpb.ArtifactKindService, service.GetTarget(), service) + resultID := toolargs.ArtifactResultID("gogo", toolpb.ArtifactKindService, service.GetTarget(), service) + c.EmitArtifactResultCtx(ctx, resultID, "gogo", toolpb.ArtifactKindService, service.GetTarget(), service) } } for _, probe := range result.Spray { if probe != nil { - c.EmitArtifactCtx(ctx, "spray", toolpb.ArtifactKindWeb, probe.UrlString, probe) + resultID := toolargs.ArtifactResultID("spray", toolpb.ArtifactKindWeb, probe.UrlString, probe) + c.EmitArtifactResultCtx(ctx, resultID, "spray", toolpb.ArtifactKindWeb, probe.UrlString, probe) } } + for _, artifact := range result.Artifacts { + c.EmitArtifactResultCtx( + ctx, + artifact.ResultID, + artifact.Tool, + artifact.Kind, + artifact.Target, + artifact.Data, + ) + } for i := range result.Loots { loot := result.Loots[i] - kind := loot.Kind - if kind == "" { - kind = toolpb.ArtifactKindVuln + resultID, _ := loot.Data["result_id"].(string) + tool, _ := loot.Data["artifact_tool"].(string) + if resultID == "" || tool == "" { + c.Logger.Warnf("skip unbound %s loot for %s", loot.Kind, loot.Target) + continue } - c.EmitArtifactCtx(ctx, "scan", kind, loot.Target, &loot) - } - for i := range result.Errors { - scanErr := result.Errors[i] - c.EmitArtifactCtx(ctx, "scan", toolpb.ArtifactKindError, scanErr.Source, &scanErr) + verificationStatus, _ := loot.Data["verification_status"].(string) + c.EmitLootCtx( + ctx, + resultID, + tool, + loot.Kind, + loot.Target, + loot.Priority, + loot.Description, + verificationStatus, + loot.Tags, + ) } - c.EmitArtifactCtx(ctx, "scan", toolpb.ArtifactKindSummary, "", &result.Summary) } var scanFileFlags = map[string]bool{ diff --git a/tools/scan/command_test.go b/tools/scan/command_test.go index 5ad4a9b2..7959fc2a 100644 --- a/tools/scan/command_test.go +++ b/tools/scan/command_test.go @@ -106,6 +106,29 @@ func TestScanProfilesAssembleCapabilities(t *testing.T) { } } +func TestBroadPOCFlagDerivesUnfingerprintedTarget(t *testing.T) { + profile, err := profileForFlags(flags{Mode: scanModeQuick, BroadPOC: true}) + if err != nil { + t.Fatalf("profileForFlags() error = %v", err) + } + if !profile.AllowBroadPOC { + t.Fatal("profile dropped --broad-poc") + } + + result := &parsers.SprayResult{ + IsValid: true, + UrlString: "http://127.0.0.1:8080", + Status: http.StatusOK, + } + var events []event + deriveWebProbeResult(profile, capSprayCheck, result, "", func(event event) { + events = append(events, event) + }) + if !hasTargetKind(events, targetPOC) { + t.Fatalf("derived events missing broad poc target: %#v", events) + } +} + func TestScanOptionsResolveCredentialFlags(t *testing.T) { flags := flags{ Users: []string{"root", "admin"}, @@ -1673,19 +1696,65 @@ func TestEmitStructuredDataPublishesScannerFacts(t *testing.T) { }}, }) - if len(events) != 3 { - t.Fatalf("events = %d, want 3: %#v", len(events), events) + if len(events) != 2 { + t.Fatalf("events = %d, want 2: %#v", len(events), events) } wants := []struct{ tool, kind string }{ - {"gogo", toolpb.ArtifactKindService}, {"spray", toolpb.ArtifactKindWeb}, {"scan", toolpb.ArtifactKindSummary}, + {"gogo", toolpb.ArtifactKindService}, {"spray", toolpb.ArtifactKindWeb}, } for index, event := range events { artifact := new(toolpb.Artifact) if event.GetExtension() == nil || event.GetExtension().UnmarshalTo(artifact) != nil { t.Fatalf("artifact event = %#v", event) } - if artifact.Tool != wants[index].tool || artifact.Kind != wants[index].kind || artifact.CallId != "scan-call-1" { + if artifact.Tool != wants[index].tool || artifact.Kind != wants[index].kind || artifact.CallId != "scan-call-1" || artifact.ResultId == "" { t.Fatalf("artifact = %#v, want %#v", artifact, wants[index]) } } } + +func TestEmitStructuredDataPublishesNativeArtifactAndLoot(t *testing.T) { + bus := eventbus.New[*aop.Event]() + cmd := New(&engine.Set{}, WithEvents(bus)) + var events []*aop.Event + unsub := bus.Subscribe(func(event *aop.Event) { events = append(events, event) }) + defer unsub() + + ctx := coretool.ContextWithInvocation(context.Background(), coretool.Invocation{ + CallID: "scan-call-1", SessionID: "scan-session", TurnID: "scan-turn", Emitter: "scan", + }) + record := &sdktypes.TemplateResult{ + Target: "http://127.0.0.1:5000", TemplateID: "test-rce", Severity: "critical", Matched: true, + Request: "GET / HTTP/1.1", Response: "HTTP/1.1 200 OK", + } + cmd.emitStructuredData(ctx, &output.ScanResult{ + Artifacts: []output.ArtifactResult{{ + ResultID: "result-1", Tool: "neutron", Kind: toolpb.ArtifactKindVuln, + Target: record.Target, Data: record, + }}, + Loots: []output.Loot{{ + Kind: output.LootVuln, Target: record.Target, Priority: "critical", + Tags: []string{"rce"}, Data: map[string]any{ + "result_id": "result-1", "artifact_tool": "neutron", "verification_status": "confirmed", + }, + }}, + }) + + if len(events) != 2 { + t.Fatalf("events = %d, want artifact + loot", len(events)) + } + artifact := new(toolpb.Artifact) + if err := events[0].GetExtension().UnmarshalTo(artifact); err != nil { + t.Fatal(err) + } + if artifact.Tool != "neutron" || artifact.ResultId != "result-1" { + t.Fatalf("artifact = %#v", artifact) + } + loot := new(toolpb.Loot) + if err := events[1].GetExtension().UnmarshalTo(loot); err != nil { + t.Fatal(err) + } + if loot.Tool != "neutron" || loot.ResultId != artifact.ResultId || loot.VerificationStatus != "confirmed" { + t.Fatalf("loot = %#v", loot) + } +} diff --git a/tools/scan/event.go b/tools/scan/event.go index b35a1c09..9c58f510 100644 --- a/tools/scan/event.go +++ b/tools/scan/event.go @@ -23,13 +23,14 @@ const ( var statsEventSeq uint64 type event struct { - Kind eventKind - Source string - Raw string - Target target - Loot *output.Loot - Error errorEvent - Stats sdktypes.Stats + Kind eventKind + Source string + Raw string + Target target + Loot *output.Loot + Artifact *output.ArtifactResult + Error errorEvent + Stats sdktypes.Stats } func targetEvent(source, raw string, target target) event { @@ -43,6 +44,19 @@ func lootEvent(source string, loot output.Loot) event { return event{Kind: eventLoot, Source: source, Loot: &loot} } +func bindLoot(loot output.Loot, resultID, tool string) output.Loot { + if loot.Data == nil { + loot.Data = make(map[string]any) + } + loot.Data["result_id"] = resultID + loot.Data["artifact_tool"] = tool + return loot +} + +func artifactLootEvent(source string, loot output.Loot, artifact output.ArtifactResult) event { + return event{Kind: eventLoot, Source: source, Loot: &loot, Artifact: &artifact} +} + func errorEventOf(source, message string) event { return event{Kind: eventError, Source: source, Error: errorEvent{Message: message}} } diff --git a/tools/scan/scan_options.go b/tools/scan/scan_options.go index 3dcc5f80..0603ba38 100644 --- a/tools/scan/scan_options.go +++ b/tools/scan/scan_options.go @@ -102,6 +102,15 @@ type profile struct { AllowBroadPOC bool } +func profileForFlags(flags flags) (profile, error) { + profile, err := profileForMode(flags.Mode) + if err != nil { + return profile, err + } + profile.AllowBroadPOC = flags.BroadPOC + return profile, nil +} + func profileForMode(mode string) (profile, error) { mode = strings.ToLower(strings.TrimSpace(mode)) if mode == "" { diff --git a/tools/scan/structured.go b/tools/scan/structured.go index 4771b70c..aeadf960 100644 --- a/tools/scan/structured.go +++ b/tools/scan/structured.go @@ -39,6 +39,7 @@ func (c *collector) StructuredResult() *output.ScanResult { } result.Spray = append(result.Spray, item.Result) } + result.Artifacts = append(result.Artifacts, c.artifacts...) result.Loots = append(result.Loots, c.loots...) for _, message := range c.errors { result.Errors = append(result.Errors, output.Error{Message: message}) diff --git a/tools/toolargs/base.go b/tools/toolargs/base.go index 2e561635..9f243ce9 100644 --- a/tools/toolargs/base.go +++ b/tools/toolargs/base.go @@ -2,7 +2,9 @@ package toolargs import ( "context" + "crypto/sha256" "encoding/json" + "fmt" "time" aop "github.com/chainreactors/aiscan/aop" @@ -32,6 +34,18 @@ func (b *Base) InitLogger(logger telemetry.Logger) { } func (b *Base) EmitArtifactCtx(ctx context.Context, tool, kind, target string, data any) { + b.EmitArtifactResultCtx(ctx, ArtifactResultID(tool, kind, target, data), tool, kind, target, data) +} + +// ArtifactResultID returns a stable identity for one scanner-native record. +// The same value is carried by its aop.tool.Loot marker. +func ArtifactResultID(tool, kind, target string, data any) string { + raw, _ := json.Marshal(data) + digest := sha256.Sum256(append([]byte(tool+"\x00"+kind+"\x00"+target+"\x00"), raw...)) + return fmt.Sprintf("%x", digest[:16]) +} + +func (b *Base) EmitArtifactResultCtx(ctx context.Context, resultID, tool, kind, target string, data any) { if b.Events == nil || data == nil { return } @@ -44,6 +58,7 @@ func (b *Base) EmitArtifactCtx(ctx context.Context, tool, kind, target string, d artifact := &toolpb.Artifact{ Tool: tool, Kind: kind, Target: target, Data: raw, MediaType: aop.JSONMediaType, Timestamp: timestamppb.New(time.Now()), CallId: invocation.CallID, + ResultId: resultID, } extension, err := anypb.New(artifact) if err != nil { @@ -59,3 +74,38 @@ func (b *Base) EmitArtifactCtx(ctx context.Context, tool, kind, target string, d Payload: &aop.Event_Extension{Extension: extension}, }) } + +func (b *Base) EmitLootCtx( + ctx context.Context, + resultID, tool, kind, target, priority, description, verificationStatus string, + tags []string, +) { + if b.Events == nil || resultID == "" { + return + } + invocation := coretool.InvocationFromContext(ctx) + loot := &toolpb.Loot{ + ResultId: resultID, + Tool: tool, + Kind: kind, + Target: target, + Priority: priority, + Tags: append([]string(nil), tags...), + Description: description, + VerificationStatus: verificationStatus, + CallId: invocation.CallID, + } + extension, err := anypb.New(loot) + if err != nil { + b.Logger.Warnf("encode %s loot: %s", tool, err) + return + } + emitter := invocation.Emitter + if emitter == "" { + emitter = tool + } + b.Events.Emit(&aop.Event{ + SessionId: invocation.SessionID, TurnId: invocation.TurnID, Emitter: emitter, + Payload: &aop.Event_Extension{Extension: extension}, + }) +} diff --git a/web/frontend/cyber-ui b/web/frontend/cyber-ui index de41e92b..236facf9 160000 --- a/web/frontend/cyber-ui +++ b/web/frontend/cyber-ui @@ -1 +1 @@ -Subproject commit de41e92bd43b554dab39be7de9e6ac7c36289418 +Subproject commit 236facf99fbe934c6b9131c9d5df444bae18d5c0