package com.sensepost.yeti.controllers; import java.io.FileNotFoundException; import java.io.IOException; import java.util.ArrayList; import java.util.HashMap; import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Queue; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.logging.Level; import java.util.logging.Logger; import javax.swing.SwingWorker; import javax.swing.table.DefaultTableModel; import com.sensepost.yeti.models.TldModel; import com.sensepost.yeti.results.DomainResult; import com.sensepost.yeti.common.ConfigSettings; import com.sensepost.yeti.common.Log; import com.sensepost.yeti.common.NetworkTools; import com.sensepost.yeti.gui.DisplayResultIFace; import com.sensepost.yeti.gui.TLDInit; import com.sensepost.yeti.gui.TldDisplayResults; import com.sensepost.yeti.helpers.TldExpandHelper; import com.sensepost.yeti.persistence.DataStore; import org.xbill.DNS.CNAMERecord; import org.xbill.DNS.Lookup; import org.xbill.DNS.Record; import org.xbill.DNS.SOARecord; import org.xbill.DNS.TextParseException; import org.xbill.DNS.Type; /** * * @author willemm */ public class TldController extends BaseController { private List<String> tlds; private final List<TldWorker> workers; private final Queue<String> urls = new LinkedList<>(); private final int retryCount = ConfigSettings.getTldRetryCount(); private final ExecutorService tldExecutor; public TldController(WorkflowController parent, TldModel dataModel) { setName("TldExpand"); setWorkflowController(parent); DisplayResultIFace displayResults = new TldDisplayResults(); displayResults.setModel((DefaultTableModel) dataModel); setDisplayResults(displayResults); workers = new ArrayList<>(); tldExecutor = Executors.newFixedThreadPool(ConfigSettings.getTldThreadCount()); } /** * Retrieves and removes the head of this queue, or returns null if this * queue is empty. * * @return Head of the domains queue */ private synchronized String popTarget() { return (String) urls.poll(); } @Override public boolean startJob() { TLDInit frmInit = new TLDInit(); frmInit.setVisible(true); if (!frmInit.cancelled) { try { List<String> rootNameList = frmInit.getRootTerms(); String tldFileLocation = frmInit.getTLDLocation(); // the tld list file's text tlds = TldExpandHelper.getTldList(tldFileLocation); // Get all the URLs to be tested for (String rootName : rootNameList) { for (String tld : tlds) { urls.add(rootName + tld); } } setRequestsCompleted(0); setTotalRequests(urls.size()); String domain = popTarget(); while (domain != null && !domain.isEmpty()) { if (getCancelledValue()) { break; } TldWorker tldExpandRunner = new TldWorker(domain); tldExecutor.execute(tldExpandRunner); workers.add(tldExpandRunner); domain = popTarget(); } } catch (FileNotFoundException ex) { Logger.getLogger(TldController.class.getName()).log(Level.SEVERE, null, ex); } catch (IOException ex) { Logger.getLogger(TldController.class.getName()).log(Level.SEVERE, null, ex); } } return !frmInit.cancelled; } private synchronized void workerDone(TldWorker worker) { workers.remove(worker); updateProgress(1); if (workers.isEmpty()) { if (getCancelledValue()) { cancelProgress(); } finishProgress(); } } @Override protected void addResult(Object result) { ((TldModel) (getDisplayResults().getModel())).addRow(result); } @Override protected void addResults(ArrayList<Object> results) { for (Object obj : results) { ((TldModel) (getDisplayResults().getModel())).addRow(obj); } } @Override public void saveData() { DataStore.addDomains(((TldModel) (getDisplayResults().getModel())).getData()); } class TldWorker extends SwingWorker<Object, Object> { private final String domain; public TldWorker(String domain) { this.domain = domain; } @Override protected String doInBackground() { try { expandDomain(domain); } catch (TextParseException ex) { Logger.getLogger(TldController.class.getName() + ":TldWorker -> doInBackground").log(Level.SEVERE, null, ex); } return "Done"; } public void expandDomain(String domainToCheck) throws TextParseException { String domainName = null; // Check for domain name alias - CNAME Record[] recs = new Lookup(domainToCheck, Type.CNAME).run(); if (recs != null && recs.length != 0) { domainName = ((CNAMERecord) recs[0]).getName().canonicalize().toString(true); Log.debug("Found: " + domainName + "CNAME rec: " + domainName); } // Now get the SOA record that would signify a domain exists recs = new Lookup(domainToCheck, Type.SOA).run(); for (int idx = 0; idx < retryCount; idx++) { if (recs != null) { if (domainName == null) { domainName = ((SOARecord) recs[0]).getName().canonicalize().toString(true); Log.debug("Found: " + domainName + " SOA rec: " + domainName); } DomainResult newDomain = new DomainResult(domainName); newDomain.setNameServer(((SOARecord) recs[0]).getHost().toString(true)); newDomain.setAdminName(((SOARecord) recs[0]).getAdmin().toString(true)); String name = domainToCheck.split("\\.", 2)[0]; String tld = domainToCheck.split("\\.", 2)[1]; newDomain.setRegistrant(NetworkTools.getHostNameWhoisResult(name, tld, true)); Map<String, String> attrs = new HashMap<>(); attrs.put(DataStore.DNS_RECORD, DataStore.SOA); newDomain.setAttributes(attrs); addResult(newDomain); break; } } } @Override protected void done() { Log.debug("Done with domain expand for: " + domain); workerDone(this); } } }