Examples with TeeOutputStream com.unboundid.util.TeeOutputStream used on opensource projects

Example 1 with TeeOutputStream

use of com.unboundid.util.TeeOutputStream in project ldapsdk by pingidentity.

the class LDAPSearch method doToolProcessing.

/**
 * {@inheritDoc}
 */
@Override()
@NotNull()
public ResultCode doToolProcessing() {
    // If we should encrypt the output, then get the encryption passphrase.
    if (encryptOutput.isPresent()) {
        if (encryptionPassphraseFile.isPresent()) {
            try {
                encryptionPassphrase = ToolUtils.readEncryptionPassphraseFromFile(encryptionPassphraseFile.getValue());
            } catch (final LDAPException e) {
                Debug.debugException(e);
                wrapErr(0, WRAP_COLUMN, e.getMessage());
                return e.getResultCode();
            }
        } else {
            try {
                encryptionPassphrase = ToolUtils.promptForEncryptionPassphrase(false, true, getOut(), getErr());
            } catch (final LDAPException e) {
                Debug.debugException(e);
                wrapErr(0, WRAP_COLUMN, e.getMessage());
                return e.getResultCode();
            }
        }
    }
    // the header to standard output.
    if (outputFile.isPresent()) {
        if (!separateOutputFilePerSearch.isPresent()) {
            try {
                OutputStream s = new FileOutputStream(outputFile.getValue());
                if (encryptOutput.isPresent()) {
                    s = new PassphraseEncryptedOutputStream(encryptionPassphrase, s);
                }
                if (compressOutput.isPresent()) {
                    s = new GZIPOutputStream(s);
                }
                if (teeResultsToStandardOut.isPresent()) {
                    outStream = new PrintStream(new TeeOutputStream(s, getOut()));
                } else {
                    outStream = new PrintStream(s);
                }
                resultWriter.updateOutputStream(outStream);
                errStream = outStream;
            } catch (final Exception e) {
                Debug.debugException(e);
                wrapErr(0, WRAP_COLUMN, ERR_LDAPSEARCH_CANNOT_OPEN_OUTPUT_FILE.get(outputFile.getValue().getAbsolutePath(), StaticUtils.getExceptionMessage(e)));
                return ResultCode.LOCAL_ERROR;
            }
            resultWriter.writeHeader();
        }
    } else {
        resultWriter.writeHeader();
    }
    // Examine the arguments to determine the sets of controls to use for each
    // type of request.
    final List<Control> searchControls = getSearchControls();
    // If appropriate, ensure that any search result entries that include
    // base64-encoded attribute values will also include comments that attempt
    // to provide a human-readable representation of that value.
    final boolean originalCommentAboutBase64EncodedValues = LDIFWriter.commentAboutBase64EncodedValues();
    LDIFWriter.setCommentAboutBase64EncodedValues(!suppressBase64EncodedValueComments.isPresent());
    LDAPConnectionPool pool = null;
    try {
        // directory server.
        if (!dryRun.isPresent()) {
            try {
                final StartAdministrativeSessionPostConnectProcessor p;
                if (useAdministrativeSession.isPresent()) {
                    p = new StartAdministrativeSessionPostConnectProcessor(new StartAdministrativeSessionExtendedRequest(getToolName(), true));
                } else {
                    p = null;
                }
                pool = getConnectionPool(1, 1, 0, p, null, true, new ReportBindResultLDAPConnectionPoolHealthCheck(this, true, false));
            } catch (final LDAPException le) {
                // This shouldn't happen since the pool won't throw an exception if an
                // attempt to create an initial connection fails.
                Debug.debugException(le);
                commentToErr(ERR_LDAPSEARCH_CANNOT_CREATE_CONNECTION_POOL.get(StaticUtils.getExceptionMessage(le)));
                return le.getResultCode();
            }
            if (retryFailedOperations.isPresent()) {
                pool.setRetryFailedOperationsDueToInvalidConnections(true);
            }
        }
        // If appropriate, create a rate limiter.
        final FixedRateBarrier rateLimiter;
        if (ratePerSecond.isPresent()) {
            rateLimiter = new FixedRateBarrier(1000L, ratePerSecond.getValue());
        } else {
            rateLimiter = null;
        }
        // requests from those URLs.
        if (ldapURLFile.isPresent()) {
            return searchWithLDAPURLs(pool, rateLimiter, searchControls);
        }
        // Get the set of requested attributes, as a combination of the
        // requestedAttribute argument values and any trailing arguments.
        final ArrayList<String> attrList = new ArrayList<>(10);
        if (requestedAttribute.isPresent()) {
            attrList.addAll(requestedAttribute.getValues());
        }
        final List<String> trailingArgs = parser.getTrailingArguments();
        if (!trailingArgs.isEmpty()) {
            final Iterator<String> trailingArgIterator = trailingArgs.iterator();
            if (!(filter.isPresent() || filterFile.isPresent())) {
                trailingArgIterator.next();
            }
            while (trailingArgIterator.hasNext()) {
                attrList.add(trailingArgIterator.next());
            }
        }
        final String[] attributes = new String[attrList.size()];
        attrList.toArray(attributes);
        // If either or both the filter or filterFile arguments are provided, then
        // use them to get the filters to process.  Otherwise, the first trailing
        // argument should be a filter.
        ResultCode resultCode = ResultCode.SUCCESS;
        if (filter.isPresent() || filterFile.isPresent()) {
            if (filter.isPresent()) {
                for (final Filter f : filter.getValues()) {
                    final ResultCode rc = searchWithFilter(pool, f, attributes, rateLimiter, searchControls);
                    if (rc != ResultCode.SUCCESS) {
                        if (resultCode == ResultCode.SUCCESS) {
                            resultCode = rc;
                        }
                        if (!continueOnError.isPresent()) {
                            return resultCode;
                        }
                    }
                }
            }
            if (filterFile.isPresent()) {
                final ResultCode rc = searchWithFilterFile(pool, attributes, rateLimiter, searchControls);
                if (rc != ResultCode.SUCCESS) {
                    if (resultCode == ResultCode.SUCCESS) {
                        resultCode = rc;
                    }
                    if (!continueOnError.isPresent()) {
                        return resultCode;
                    }
                }
            }
        } else {
            final Filter f;
            try {
                final String filterStr = parser.getTrailingArguments().iterator().next();
                f = Filter.create(filterStr);
            } catch (final LDAPException le) {
                // This should never happen.
                Debug.debugException(le);
                displayResult(le.toLDAPResult());
                return le.getResultCode();
            }
            resultCode = searchWithFilter(pool, f, attributes, rateLimiter, searchControls);
        }
        return resultCode;
    } finally {
        if (pool != null) {
            try {
                pool.close();
            } catch (final Exception e) {
                Debug.debugException(e);
            }
        }
        if (outStream != null) {
            try {
                outStream.close();
                outStream = null;
            } catch (final Exception e) {
                Debug.debugException(e);
            }
        }
        if (errStream != null) {
            try {
                errStream.close();
                errStream = null;
            } catch (final Exception e) {
                Debug.debugException(e);
            }
        }
        LDIFWriter.setCommentAboutBase64EncodedValues(originalCommentAboutBase64EncodedValues);
    }
}

Example 2 with TeeOutputStream

use of com.unboundid.util.TeeOutputStream in project ldapsdk by pingidentity.

the class LDAPSearch method doSearch.

/**
 * Processes a search with the provided information.
 *
 * @param  pool            The connection pool to use to communicate with the
 *                         directory server.
 * @param  searchRequest   The search request to process.
 * @param  rateLimiter     An optional fixed-rate barrier that can be used for
 *                         request rate limiting.
 * @param  searchControls  The set of controls to include in search requests.
 *
 * @return  A result code indicating the result of the processing.
 */
@NotNull()
private ResultCode doSearch(@NotNull final LDAPConnectionPool pool, @NotNull final SearchRequest searchRequest, @Nullable final FixedRateBarrier rateLimiter, @NotNull final List<Control> searchControls) {
    if (separateOutputFilePerSearch.isPresent()) {
        try {
            final String path = outputFile.getValue().getAbsolutePath() + '.' + outputFileCounter.getAndIncrement();
            OutputStream s = new FileOutputStream(path);
            if (encryptOutput.isPresent()) {
                s = new PassphraseEncryptedOutputStream(encryptionPassphrase, s);
            }
            if (compressOutput.isPresent()) {
                s = new GZIPOutputStream(s);
            }
            if (teeResultsToStandardOut.isPresent()) {
                outStream = new PrintStream(new TeeOutputStream(s, getOut()));
            } else {
                outStream = new PrintStream(s);
            }
            resultWriter.updateOutputStream(outStream);
            errStream = outStream;
        } catch (final Exception e) {
            Debug.debugException(e);
            wrapErr(0, WRAP_COLUMN, ERR_LDAPSEARCH_CANNOT_OPEN_OUTPUT_FILE.get(outputFile.getValue().getAbsolutePath(), StaticUtils.getExceptionMessage(e)));
            return ResultCode.LOCAL_ERROR;
        }
        resultWriter.writeHeader();
    }
    try {
        if (rateLimiter != null) {
            rateLimiter.await();
        }
        ASN1OctetString pagedResultsCookie = null;
        boolean multiplePages = false;
        long totalEntries = 0;
        long totalReferences = 0;
        SearchResult searchResult;
        try {
            while (true) {
                searchRequest.setControls(searchControls);
                if (simplePageSize.isPresent()) {
                    searchRequest.addControl(new SimplePagedResultsControl(simplePageSize.getValue(), pagedResultsCookie));
                }
                if (dryRun.isPresent()) {
                    searchResult = new SearchResult(-1, ResultCode.SUCCESS, INFO_LDAPSEARCH_DRY_RUN_REQUEST_NOT_SENT.get(dryRun.getIdentifierString(), String.valueOf(searchRequest)), null, null, 0, 0, null);
                    break;
                } else {
                    if (!terse.isPresent()) {
                        if (verbose.isPresent() || persistentSearch.isPresent() || filterFile.isPresent() || ldapURLFile.isPresent() || (filter.isPresent() && (filter.getNumOccurrences() > 1))) {
                            commentToOut(INFO_LDAPSEARCH_SENDING_SEARCH_REQUEST.get(String.valueOf(searchRequest)));
                        }
                    }
                    searchResult = pool.search(searchRequest);
                }
                if (searchResult.getEntryCount() > 0) {
                    totalEntries += searchResult.getEntryCount();
                }
                if (searchResult.getReferenceCount() > 0) {
                    totalReferences += searchResult.getReferenceCount();
                }
                if (simplePageSize.isPresent()) {
                    final SimplePagedResultsControl pagedResultsControl;
                    try {
                        pagedResultsControl = SimplePagedResultsControl.get(searchResult);
                        if (pagedResultsControl == null) {
                            throw new LDAPSearchException(new SearchResult(searchResult.getMessageID(), ResultCode.CONTROL_NOT_FOUND, ERR_LDAPSEARCH_MISSING_PAGED_RESULTS_RESPONSE_CONTROL.get(), searchResult.getMatchedDN(), searchResult.getReferralURLs(), searchResult.getSearchEntries(), searchResult.getSearchReferences(), searchResult.getEntryCount(), searchResult.getReferenceCount(), searchResult.getResponseControls()));
                        }
                        if (pagedResultsControl.moreResultsToReturn()) {
                            if (verbose.isPresent()) {
                                commentToOut(INFO_LDAPSEARCH_INTERMEDIATE_PAGED_SEARCH_RESULT.get());
                                displayResult(searchResult);
                            }
                            multiplePages = true;
                            pagedResultsCookie = pagedResultsControl.getCookie();
                        } else {
                            break;
                        }
                    } catch (final LDAPException le) {
                        Debug.debugException(le);
                        throw new LDAPSearchException(new SearchResult(searchResult.getMessageID(), ResultCode.CONTROL_NOT_FOUND, ERR_LDAPSEARCH_CANNOT_DECODE_PAGED_RESULTS_RESPONSE_CONTROL.get(StaticUtils.getExceptionMessage(le)), searchResult.getMatchedDN(), searchResult.getReferralURLs(), searchResult.getSearchEntries(), searchResult.getSearchReferences(), searchResult.getEntryCount(), searchResult.getReferenceCount(), searchResult.getResponseControls()));
                    }
                } else {
                    break;
                }
            }
        } catch (final LDAPSearchException lse) {
            Debug.debugException(lse);
            searchResult = lse.toLDAPResult();
            if (searchResult.getEntryCount() > 0) {
                totalEntries += searchResult.getEntryCount();
            }
            if (searchResult.getReferenceCount() > 0) {
                totalReferences += searchResult.getReferenceCount();
            }
        }
        if ((searchResult.getResultCode() != ResultCode.SUCCESS) || (searchResult.getDiagnosticMessage() != null) || (!terse.isPresent())) {
            displayResult(searchResult);
        }
        if (multiplePages && (!terse.isPresent())) {
            commentToOut(INFO_LDAPSEARCH_TOTAL_SEARCH_ENTRIES.get(totalEntries));
            if (totalReferences > 0) {
                commentToOut(INFO_LDAPSEARCH_TOTAL_SEARCH_REFERENCES.get(totalReferences));
            }
        }
        if (countEntries.isPresent()) {
            return ResultCode.valueOf((int) Math.min(totalEntries, 255));
        } else if (requireMatch.isPresent() && (totalEntries == 0)) {
            return ResultCode.NO_RESULTS_RETURNED;
        } else {
            return searchResult.getResultCode();
        }
    } finally {
        if (separateOutputFilePerSearch.isPresent()) {
            try {
                outStream.close();
            } catch (final Exception e) {
                Debug.debugException(e);
            }
            outStream = null;
            errStream = null;
        }
    }
}