HostAddressManager.java

Go to the documentation of this file.

/*
 * Sleuth Kit Data Model
 *
 * Copyright 2021 Basis Technology Corp.
 * Contact: carrier <at> sleuthkit <dot> org
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.sleuthkit.datamodel;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.regex.Pattern;
import org.sleuthkit.datamodel.SleuthkitCase.CaseDbConnection;
import org.sleuthkit.datamodel.HostAddress.HostAddressType;

public class HostAddressManager {

        private static final Logger LOGGER = Logger.getLogger(HostAddressManager.class.getName());

        private final SleuthkitCase db;

        private final Cache<Long, Byte> recentHostNameAndIpMappingCache = CacheBuilder.newBuilder().maximumSize(200000).build();

        private final Cache<String, HostAddress> recentHostAddressCache = CacheBuilder.newBuilder().maximumSize(200000).build();

        private final Cache<String, Boolean> hostAddressUsageCache = CacheBuilder.newBuilder().maximumSize(200000).build();

        HostAddressManager(SleuthkitCase skCase) {
                this.db = skCase;
        }

        public Optional<HostAddress> getHostAddress(HostAddress.HostAddressType type, String address) throws TskCoreException {
                
                db.acquireSingleUserCaseReadLock();
                try (CaseDbConnection connection = this.db.getConnection()) {
                        return HostAddressManager.this.getHostAddress(type, address, connection);
                } finally {
                        db.releaseSingleUserCaseReadLock();
                }
        }

        private Optional<HostAddress> getHostAddress(HostAddress.HostAddressType type, String address, CaseDbConnection connection) throws TskCoreException {
                
                HostAddress hostAddress = recentHostAddressCache.getIfPresent(createRecentHostAddressKey(type, address));
                if (Objects.nonNull(hostAddress)) {
                        return Optional.of(hostAddress);
                }
                HostAddress.HostAddressType addressType = type;
                if (type.equals(HostAddress.HostAddressType.DNS_AUTO)) {
                        addressType = getDNSType(address);
                }
                String queryString = "SELECT * FROM tsk_host_addresses"
                                + " WHERE address = ?  AND address_type = ?";                   
                try {
                        PreparedStatement query = connection.getPreparedStatement(queryString, Statement.NO_GENERATED_KEYS);
                        query.clearParameters();
                        query.setString(1, address.toLowerCase());
                        query.setInt(2, addressType.getId());
                        try (ResultSet rs = query.executeQuery()) {
                                if (!rs.next()) {
                                        return Optional.empty();        // no match found
                                } else {
                                        HostAddress newHostAddress = new HostAddress(db, rs.getLong("id"), HostAddressType.fromID(rs.getInt("address_type")), address);
                                        recentHostAddressCache.put(createRecentHostAddressKey(newHostAddress.getAddressType(), address), newHostAddress);
                                        return Optional.of(newHostAddress);                                     
                                }
                        }
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error getting host address with type = %s and address = %s", type.getName(), address), ex);
                } 
        }

        private String createRecentHostAddressKey(HostAddressType type, String address) {
                return createRecentHostAddressKey(type.getId(), address);
        }

        private String createRecentHostAddressKey(int typeId, String address) {
                return typeId + "#" + address.toLowerCase();
        }

        public HostAddress newHostAddress(HostAddress.HostAddressType type, String address) throws TskCoreException {
                db.acquireSingleUserCaseWriteLock();
                CaseDbConnection connection = this.db.getConnection();
                try {
                        return HostAddressManager.this.newHostAddress(type, address, connection);
                } catch (TskCoreException ex) {
                        // The insert may have failed because the HostAddress already exists, so
                        // try loading it from the database.
                        Optional<HostAddress> hostAddress = HostAddressManager.this.getHostAddress(type, address, connection);
                        if (hostAddress.isPresent()) {
                                return hostAddress.get();
                        }
                        throw ex;
                } finally {
                        connection.close(); 
                        db.releaseSingleUserCaseWriteLock();
                }
        }

        private HostAddress newHostAddress(HostAddress.HostAddressType type, String address, CaseDbConnection connection) throws TskCoreException {
                HostAddress.HostAddressType addressType = type;
                if (type.equals(HostAddress.HostAddressType.DNS_AUTO)) {
                        addressType = getDNSType(address);
                }
                
                try {

                        // TODO: need to get the correct parent obj id.  
                        long parentObjId = 0;
                        int objTypeId = TskData.ObjectType.HOST_ADDRESS.getObjectType();

                        long objId = db.addObject(parentObjId, objTypeId, connection);

                        String hostAddressInsertSQL = "INSERT INTO tsk_host_addresses(id, address_type, address) VALUES (?, ?, ?)"; // NON-NLS
                        PreparedStatement preparedStatement = connection.getPreparedStatement(hostAddressInsertSQL, Statement.RETURN_GENERATED_KEYS);

                        preparedStatement.clearParameters();
                        preparedStatement.setLong(1, objId);
                        preparedStatement.setInt(2, addressType.getId());
                        preparedStatement.setString(3, address.toLowerCase());

                        connection.executeUpdate(preparedStatement);
                        HostAddress hostAddress =  new HostAddress(db, objId, addressType, address);
                        recentHostAddressCache.put(createRecentHostAddressKey(addressType, address), hostAddress);
                        return hostAddress;
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error adding host address of type = %s, with address = %s", type.getName(), address), ex);
                } 
        }

        public void assignHostToAddress(Host host, HostAddress hostAddress, Long time, Content source) throws TskCoreException {

                String insertSQL = db.getInsertOrIgnoreSQL(" INTO tsk_host_address_map(host_id, addr_obj_id, source_obj_id, time) "
                                + " VALUES(?, ?, ?, ?) ");

                db.acquireSingleUserCaseWriteLock();
                try (CaseDbConnection connection = this.db.getConnection()) {

                        PreparedStatement preparedStatement = connection.getPreparedStatement(insertSQL, Statement.NO_GENERATED_KEYS);

                        preparedStatement.clearParameters();
                        preparedStatement.setLong(1, host.getHostId());
                        preparedStatement.setLong(2, hostAddress.getId());
                        preparedStatement.setLong(3, source.getId());
                        if (time != null) {
                                preparedStatement.setLong(4, time);
                        } else {
                                preparedStatement.setNull(4, java.sql.Types.BIGINT);
                        }

                        connection.executeUpdate(preparedStatement);
                } catch (SQLException ex) {
                        LOGGER.log(Level.SEVERE, null, ex);
                        throw new TskCoreException(String.format("Error adding host address mapping for host name = %s,  with address = %s", host.getName(), hostAddress.getAddress()), ex);
                } finally {
                        db.releaseSingleUserCaseWriteLock();
                }
        }

        List<HostAddress> getHostAddressesAssignedTo(Host host) throws TskCoreException {

                String queryString = "SELECT addr_obj_id FROM tsk_host_address_map "
                                + " WHERE host_id = " + host.getHostId();

                List<HostAddress> addresses = new ArrayList<>();

                db.acquireSingleUserCaseReadLock();
                try (CaseDbConnection connection = this.db.getConnection();
                                Statement s = connection.createStatement();
                                ResultSet rs = connection.executeQuery(s, queryString)) {

                        while (rs.next()) {
                                addresses.add(HostAddressManager.this.getHostAddress(rs.getLong("addr_obj_id"), connection));
                        }

                        return addresses;
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error getting host addresses for host " + host.getName()), ex);
                } finally {
                        db.releaseSingleUserCaseReadLock();
                }
        }

        public HostAddress getHostAddress(long id) throws TskCoreException {
                db.acquireSingleUserCaseReadLock();
                try (CaseDbConnection connection = this.db.getConnection()) {
                        return HostAddressManager.this.getHostAddress(id, connection);
                } finally {
                        db.releaseSingleUserCaseReadLock();
                }
        }

        private HostAddress getHostAddress(long id, CaseDbConnection connection) throws TskCoreException {
                String queryString = "SELECT * FROM tsk_host_addresses"
                                + " WHERE id = " + id;

                try (Statement s = connection.createStatement();
                                ResultSet rs = connection.executeQuery(s, queryString)) {

                        if (!rs.next()) {
                                throw new TskCoreException(String.format("No address found with id = %d", id));
                        } else {
                                long objId = rs.getLong("id");
                                int type = rs.getInt("address_type");
                                String address = rs.getString("address");
                                HostAddress hostAddress = new HostAddress(db, objId, HostAddress.HostAddressType.fromID(type), address);
                                recentHostAddressCache.put(createRecentHostAddressKey(type, address), hostAddress);
                                return hostAddress;
                        }
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error getting host address with id = %d", id), ex);
                } 
        }

        public void addHostNameAndIpMapping(HostAddress dnsNameAddress, HostAddress ipAddress, Long time, Content source) throws TskCoreException {

                db.acquireSingleUserCaseWriteLock();
                try (CaseDbConnection connection = this.db.getConnection()) {
                        addHostNameAndIpMapping(dnsNameAddress, ipAddress, time, source, connection);
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error adding host DNS address mapping for DNS name = %s, and IP address = %s", dnsNameAddress.getAddress(), ipAddress.getAddress()), ex);
                } finally {
                        db.releaseSingleUserCaseWriteLock();
                }
        }

        public void addHostNameAndIpMapping(HostAddress dnsNameAddress, HostAddress ipAddress, Long time, Content source, final SleuthkitCase.CaseDbTransaction caseDbTransaction) throws TskCoreException {

                if (Objects.isNull(caseDbTransaction)) {
                        throw new TskCoreException(String.format("Error adding host DNS address mapping for DNS name = %s, and IP address = %s, null caseDbTransaction passed to addHostNameAndIpMapping", dnsNameAddress.getAddress(), ipAddress.getAddress()));
                }
                try {
                        addHostNameAndIpMapping(dnsNameAddress, ipAddress, time, source, caseDbTransaction.getConnection());
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error adding host DNS address mapping for DNS name = %s, and IP address = %s", dnsNameAddress.getAddress(), ipAddress.getAddress()), ex);
                }
        }

        private void addHostNameAndIpMapping(HostAddress dnsNameAddress, HostAddress ipAddress, Long time, Content source, final CaseDbConnection connection) throws SQLException, TskCoreException {

                if (dnsNameAddress.getAddressType() != HostAddress.HostAddressType.HOSTNAME) {
                        throw new TskCoreException("IllegalArguments passed to addHostNameAndIpMapping: A host name address is expected.");
                }
                if ((ipAddress.getAddressType() != HostAddress.HostAddressType.IPV4) && (ipAddress.getAddressType() != HostAddress.HostAddressType.IPV6)) {
                        throw new TskCoreException("IllegalArguments passed to addHostNameAndIpMapping:An IPv4/IPv6 address is expected.");
                }
                if (Objects.isNull(connection)) {
                        throw new TskCoreException("IllegalArguments passed to addHostNameAndIpMapping: null connection passed to addHostNameAndIpMapping");
                }

                String insertSQL = db.getInsertOrIgnoreSQL(" INTO tsk_host_address_dns_ip_map(dns_address_id, ip_address_id, source_obj_id, time) "
                                + " VALUES(?, ?, ?, ?) ");

                PreparedStatement preparedStatement = connection.getPreparedStatement(insertSQL, Statement.NO_GENERATED_KEYS);

                preparedStatement.clearParameters();
                preparedStatement.setLong(1, dnsNameAddress.getId());
                preparedStatement.setLong(2, ipAddress.getId());
                preparedStatement.setLong(3, source.getId());
                if (time != null) {
                        preparedStatement.setLong(4, time);
                } else {
                        preparedStatement.setNull(4, java.sql.Types.BIGINT);
                }
                connection.executeUpdate(preparedStatement);
                recentHostNameAndIpMappingCache.put(ipAddress.getId(), new Byte((byte) 1));
                recentHostNameAndIpMappingCache.put(dnsNameAddress.getId(), new Byte((byte) 1));
        }

        public boolean hostNameAndIpMappingExists(long addressObjectId) throws TskCoreException {

                Byte isPresent = recentHostNameAndIpMappingCache.getIfPresent(addressObjectId);

                if (Objects.nonNull(isPresent)) {
                        return true;
                }

                String queryString = "SELECT count(*) as mappingCount FROM tsk_host_address_dns_ip_map WHERE ip_address_id = ? OR dns_address_id = ? ";

                db.acquireSingleUserCaseReadLock();
                try (CaseDbConnection connection = this.db.getConnection();
                                PreparedStatement ps = connection.getPreparedStatement(queryString, Statement.NO_GENERATED_KEYS);) {
                        ps.clearParameters();
                        ps.setLong(1, addressObjectId);
                        ps.setLong(2, addressObjectId);
                        try (ResultSet rs = ps.executeQuery()) {
                                if (!rs.next()) {
                                        return false;
                                } else {
                                        boolean status = rs.getLong("mappingCount") > 0;
                                        if (status) {
                                                recentHostNameAndIpMappingCache.put(addressObjectId, new Byte((byte) 1));
                                        }
                                        return status;
                                }
                        }
                } catch (SQLException ex) {
                        throw new TskCoreException("Error looking up host address / Ip mapping for address = " + addressObjectId, ex);
                } finally {
                        db.releaseSingleUserCaseReadLock();
                }
        }

        public Optional<Long> hostAddressExists(HostAddress.HostAddressType type, String address) throws TskCoreException {

                HostAddress hostAddress = recentHostAddressCache.getIfPresent(createRecentHostAddressKey(type, address));
                if (Objects.nonNull(hostAddress)) {
                        return Optional.of(hostAddress.getId());
                }

                HostAddress.HostAddressType addressType = type;
                if (type.equals(HostAddress.HostAddressType.DNS_AUTO)) {
                        addressType = getDNSType(address);
                } 
                
                String queryString = "SELECT id, address_type, address FROM tsk_host_addresses"
                                + " WHERE address = ?  AND address_type = ?"; 

                db.acquireSingleUserCaseReadLock();
                try (CaseDbConnection connection = this.db.getConnection();
                                PreparedStatement query = connection.getPreparedStatement(queryString, Statement.NO_GENERATED_KEYS);) {
                        query.clearParameters();
                        query.setString(1, address.toLowerCase());
                        query.setInt(2, addressType.getId());
                        try (ResultSet rs = query.executeQuery()) {
                                if (!rs.next()) {
                                        return Optional.empty();        // no match found
                                } else {
                                        long objId = rs.getLong("id");
                                        int addrType = rs.getInt("address_type");
                                        String addr = rs.getString("address");
                                        HostAddress hostAddr = new HostAddress(db, objId, HostAddress.HostAddressType.fromID(addrType), addr);
                                        recentHostAddressCache.put(createRecentHostAddressKey(addrType, address), hostAddr);                                    
                                        return Optional.of(objId);
                                }
                        }
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error getting host address with type = %s and address = %s", type.getName(), address), ex);
                } finally {
                        db.releaseSingleUserCaseReadLock();
                }
        }

        public List<HostAddress> getIpAddress(String hostname) throws TskCoreException {
                String queryString = "SELECT ip_address_id FROM tsk_host_address_dns_ip_map as map "
                                + " JOIN tsk_host_addresses as addresses "
                                + " ON map.dns_address_id = addresses.id "
                                + " WHERE addresses.address_type = " + HostAddress.HostAddressType.HOSTNAME.getId()
                                + " AND addresses.address = ?";

                db.acquireSingleUserCaseReadLock();
                try (CaseDbConnection connection = this.db.getConnection()) {
                        List<HostAddress> IpAddresses = new ArrayList<>();
                        PreparedStatement query = connection.getPreparedStatement(queryString, Statement.NO_GENERATED_KEYS);
                        query.clearParameters();
                        query.setString(1, hostname.toLowerCase());
                        try (ResultSet rs = query.executeQuery()) {
                                while (rs.next()) {
                                        long ipAddressObjId = rs.getLong("ip_address_id");
                                        IpAddresses.add(HostAddressManager.this.getHostAddress(ipAddressObjId, connection));
                                        recentHostNameAndIpMappingCache.put(ipAddressObjId, new Byte((byte) 1));
                                }
                                return IpAddresses;
                        }
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error getting host addresses for host name: " + hostname), ex);
                } finally {
                        db.releaseSingleUserCaseReadLock();
                }
        }

        List<HostAddress> getHostNameByIp(String ipAddress) throws TskCoreException {
                String queryString = "SELECT dns_address_id FROM tsk_host_address_dns_ip_map as map "
                                + " JOIN tsk_host_addresses as addresses "
                                + " ON map.ip_address_id = addresses.id "
                                + " WHERE ( addresses.address_type = " + HostAddress.HostAddressType.IPV4.getId()
                                + " OR  addresses.address_type = " + HostAddress.HostAddressType.IPV6.getId() + ")"
                                + " AND addresses.address = ?";

                db.acquireSingleUserCaseReadLock();
                try (CaseDbConnection connection = this.db.getConnection()) {
                        List<HostAddress> dnsNames = new ArrayList<>();
                        PreparedStatement query = connection.getPreparedStatement(queryString, Statement.NO_GENERATED_KEYS);
                        query.clearParameters();
                        query.setString(1, ipAddress.toLowerCase());
                        try (ResultSet rs = query.executeQuery()) {
                                while (rs.next()) {
                                        long dnsAddressId = rs.getLong("dns_address_id");
                                        dnsNames.add(HostAddressManager.this.getHostAddress(dnsAddressId, connection));
                                        recentHostNameAndIpMappingCache.put(dnsAddressId, new Byte((byte) 1));
                                }
                                return dnsNames;
                        }
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error getting host addresses for IP address: " + ipAddress), ex);
                } finally {
                        db.releaseSingleUserCaseReadLock();
                }
        }

        public void addUsage(Content content, HostAddress hostAddress) throws TskCoreException {
                
                String key = content.getDataSource().getId() + "#" + hostAddress.getId() + "#" + content.getId();
                Boolean cachedValue = hostAddressUsageCache.getIfPresent(key);
                if (null != cachedValue) {
                        return;
                }
                
                final String insertSQL = db.getInsertOrIgnoreSQL(" INTO tsk_host_address_usage(addr_obj_id, obj_id, data_source_obj_id) "
                                + " VALUES(" + hostAddress.getId() + ", " + content.getId() + ", " + content.getDataSource().getId() + ") ");

                db.acquireSingleUserCaseWriteLock();
                try (CaseDbConnection connection = this.db.getConnection();
                                Statement s = connection.createStatement()) {
                        connection.executeUpdate(s, insertSQL);
                        hostAddressUsageCache.put(key, true);
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error associating host address %s with artifact with id %d", hostAddress.getAddress(), content.getId()), ex);
                } finally {
                        db.releaseSingleUserCaseWriteLock();
                }
        }

        private final String ADDRESS_USAGE_QUERY = "SELECT addresses.id as id, addresses.address_type as address_type, addresses.address as address "
                        + " FROM tsk_host_address_usage as usage "
                        + " JOIN tsk_host_addresses as addresses "
                        + " ON usage.addr_obj_id = addresses.id ";

        public List<HostAddress> getHostAddressesUsedByContent(Content content) throws TskCoreException {
                String queryString = ADDRESS_USAGE_QUERY
                                + " WHERE usage.obj_id = " + content.getId();

                return getHostAddressesUsed(queryString);
        }

        public List<HostAddress> getHostAddressesUsedOnDataSource(Content dataSource) throws TskCoreException {
                String queryString = ADDRESS_USAGE_QUERY
                                + " WHERE usage.data_source_obj_id = " + dataSource.getId();

                return getHostAddressesUsed(queryString);
        }

        private List<HostAddress> getHostAddressesUsed(String addressesUsedSQL) throws TskCoreException {

                List<HostAddress> addressesUsed = new ArrayList<>();

                db.acquireSingleUserCaseReadLock();
                try (CaseDbConnection connection = this.db.getConnection();
                                Statement s = connection.createStatement();
                                ResultSet rs = connection.executeQuery(s, addressesUsedSQL)) {

                        while (rs.next()) {
                                addressesUsed.add(new HostAddress(db, rs.getLong("id"), HostAddress.HostAddressType.fromID(rs.getInt("address_type")), rs.getString("address")));
                        }
                        return addressesUsed;
                } catch (SQLException ex) {
                        throw new TskCoreException(String.format("Error getting host addresses used with query string = %s", addressesUsedSQL), ex);
                } finally {
                        db.releaseSingleUserCaseReadLock();
                }
        }

        private HostAddress.HostAddressType getDNSType(String address) {
                if (isIPv4(address)) {
                        return HostAddress.HostAddressType.IPV4;
                } else if (isIPv6(address)) {
                        return HostAddress.HostAddressType.IPV6;
                } else {
                        return HostAddress.HostAddressType.HOSTNAME;
                }
        }

        private static final Pattern IPV4_PATTERN =
            Pattern.compile("^(([0-9]|[1-9][0-9]|1[0-9][0-9]|2[0-4][0-9]|25[0-5])(\\.(?!$)|$)){4}$");
        private static boolean isIPv4(String ipAddress) {
                if (ipAddress != null) {
                        return IPV4_PATTERN.matcher(ipAddress).matches();
                }
                return false;
        }

        
        private static final Pattern IPV6_STD_PATTERN = 
            Pattern.compile("^(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}$");
    private static final Pattern IPV6_HEX_COMPRESSED_PATTERN = 
            Pattern.compile("^((?:[0-9A-Fa-f]{1,4}(?::[0-9A-Fa-f]{1,4})*)?)::((?:[0-9A-Fa-f]{1,4}(?::[0-9A-Fa-f]{1,4})*)?)$");

   
    private static boolean isIPv6StdAddress(final String input) {
        return IPV6_STD_PATTERN.matcher(input).matches();
    }
    private static boolean isIPv6HexCompressedAddress(final String input) {
        return IPV6_HEX_COMPRESSED_PATTERN.matcher(input).matches();
    }
        
        private static boolean isIPv6(String ipAddress) {
        
                if (ipAddress != null) {
                         return isIPv6StdAddress(ipAddress) || isIPv6HexCompressedAddress(ipAddress);
                }

                return false;
        }
}