com.lightboxtechnologies.nsrl.HashLoader.java Source code

Introduction

Here is the source code for com.lightboxtechnologies.nsrl.HashLoader.java

Source

/*
   Copyright 2011, Lightbox Technologies, Inc
    
   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 com.lightboxtechnologies.nsrl;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.hbase.KeyValue;
import org.apache.hadoop.hbase.io.ImmutableBytesWritable;
import org.apache.hadoop.hbase.mapreduce.HFileOutputFormat;
import org.apache.hadoop.hbase.mapreduce.KeyValueSortReducer;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.Mapper;
import org.apache.hadoop.mapreduce.lib.input.TextInputFormat;
import org.apache.hadoop.util.GenericOptionsParser;

import org.sleuthkit.hadoop.core.SKJobFactory;

/**
 * An MR job to load NSRL hash data to HBase.
 *
 * @author Joel Uckelman
 */
public class HashLoader {
    static class HashLoaderMapper extends Mapper<LongWritable, Text, ImmutableBytesWritable, KeyValue> {

        private final LineTokenizer tok = new OpenCSVLineTokenizer();

        private final ErrorConsumer err = new ErrorConsumer() {
            public void consume(BadDataException e, long linenum) {
                System.err.println("malformed record, line " + linenum);
                e.printStackTrace();
            }
        };

        private final HashRecordProcessor h_proc = new HashRecordProcessor();

        private HashLoaderHelper hlh;

        @Override
        protected void setup(Context context) throws IOException, InterruptedException {
            super.setup(context);

            final Configuration conf = context.getConfiguration();

            long timestamp;
            // ensure that all mappers have the same timestamp
            try {
                timestamp = Long.parseLong(conf.get("timestamp"));
            } catch (NumberFormatException e) {
                throw new RuntimeException(e);
            }

            // load data for the manufacturer, os, and product maps
            final String mfg_filename = conf.get("mfg_filename");
            final String os_filename = conf.get("os_filename");
            final String prod_filename = conf.get("prod_filename");

            final FileSystem fs = FileSystem.get(conf);

            final Map<String, MfgData> mfg = new HashMap<String, MfgData>();
            final Map<String, OSData> os = new HashMap<String, OSData>();
            final Map<Integer, List<ProdData>> prod = new HashMap<Integer, List<ProdData>>();

            SmallTableLoader.load(fs, mfg_filename, mfg, os_filename, os, prod_filename, prod, tok, err);

            hlh = new HashLoaderHelper(prod, mfg, os, timestamp);
        }

        @Override
        protected void map(LongWritable linenum, Text line, Context context)
                throws InterruptedException, IOException {
            if (line.toString().startsWith("\"SHA-1\"")) {
                // skip this line, it's the header which gives the column names
                return;
            }

            // parse the line into the hash data
            final String[] cols = tok.tokenize(line.toString());

            HashData hd = null;
            try {
                hd = h_proc.process(cols);
            } catch (BadDataException e) {
                err.consume(e, linenum.get());
                return;
            }

            // create one row keyed to each of the md5 and sha1 hashes
            hlh.writeRow(hd.md5, hd, context);
            hlh.writeRow(hd.sha1, hd, context);
        }
    }

    public static void main(String[] args) throws Exception {
        final Configuration conf = new Configuration();

        final String[] otherArgs = new GenericOptionsParser(conf, args).getRemainingArgs();

        if (otherArgs.length != 6) {
            System.err
                    .println("Usage: HashLoader <mfgfile> <osfile> <prodfile> <hashfile> <outpath> <num_reducers>");
            System.exit(2);
        }

        final String mfg_filename = otherArgs[0];
        final String os_filename = otherArgs[1];
        final String prod_filename = otherArgs[2];
        final String hash_filename = otherArgs[3];
        final String output_filename = otherArgs[4];

        conf.set("mfg_filename", mfg_filename);
        conf.set("os_filename", os_filename);
        conf.set("prod_filename", prod_filename);

        conf.setLong("timestamp", System.currentTimeMillis());
        SKJobFactory.addDependencies(conf);

        final Job job = new Job(conf, "HashLoader");
        job.setJarByClass(HashLoader.class);
        job.setMapperClass(HashLoaderMapper.class);
        job.setReducerClass(KeyValueSortReducer.class);
        job.setNumReduceTasks(Integer.parseInt(otherArgs[5]));

        job.setOutputKeyClass(ImmutableBytesWritable.class);
        job.setOutputValueClass(KeyValue.class);

        job.setInputFormatClass(TextInputFormat.class);
        job.setOutputFormatClass(HFileOutputFormat.class);

        TextInputFormat.addInputPath(job, new Path(hash_filename));
        HFileOutputFormat.setOutputPath(job, new Path(output_filename));

        System.exit(job.waitForCompletion(true) ? 0 : 1);
    }
}